@artilleryio/platform-fargate/lib/data-api/test-run-object.js

Fargate support for Artillery

const { ObjectStore } = require('../cloud/object-store');
const { getBucketName, getAccountId } = require('../util');
const getrc = require('../utils/get-rc');
const util = require('util');
const AWS = require('aws-sdk');
const debug = require('debug')('store');
const uuidv4 = require('uuid/v4');
const { createDynamoDocumentClient } = require('../utils/create-dynamo-document-client');
const getBackendId = require('../utils/get-backend-id');
const { writeInBatches } = require('./write-in-batches');
const { ConsoleOutputSerializeError } = require('../errors');

const { StorageBackendAuroraV1 } = require('./aurora-serverless');
//
// DynamoDB-backed storage for TestRuns
//

class StorageBackendDynamo {
  constructor(id, opts) {
    this.id = id;
    this.tags = [];
    this.tagsSearchPatterns = [];
  }

  static properties = {
    type: 'aws:dynamodb',
    description: 'AWS DynamoDB backend',
  };

  async init(id, opts) {
    debug('testRunObject init');

    if (!this.tableName) {
      const accountId = await getAccountId();
      this.tableName = process.env.ARTILLERY_TEST_RUNS_TABLE || `artilleryio-test-runs-${accountId}`;
      this.commonObjectsTable = `artilleryio-common-objects-${accountId}`;
    }

    const { backendId } = await getBackendId();
    this.region = backendId;
    AWS.config.update({region: this.region});
    this.ddb = createDynamoDocumentClient();

    // Large objects such as console output live in S3:
    this.objectPrefix = `data-${this.region}`;
    const bucketName = await getBucketName();
    this.objectStore = new ObjectStore({ backend: 'aws', bucket: bucketName });

    return this;
  }

  async setChecks() {
    return this; // not implemented
  }

  async setEndedAt(ts) {
    debug('testRunObject setEndedAt');
    // Don't need this on Dynamo
    return this;
  }

  async setTagsSearchStrings(tags = [], additionalAttributes = {}) {
    const attrs = Object.entries(additionalAttributes).length > 0 ?
      { ...additionalAttributes, createdTime: this.createdTime } : additionalAttributes

    if (this.tagsSearchPatterns.length === 0) {
      this.tagsSearchPatterns = this._generateTagCombinations(
        tags
          .map((x) => x.name + ':' + x.value)
          .sort()
      )
    }

    await this._writeInBatches(this.tagsSearchPatterns.map((x) => {
      const r = {
        Item: {
          testRunId: this.id,
          kind: `tags#${x}`,
          valueType: typeof x,
          ...attrs,
        },
      };
      r.Item[`${typeof x.value}Value`] = x.value;
      return r;
    }));
  }

  async setTags(testRunTags) {
    debug('testRunObject setTags');

    this.tags = testRunTags;
    //
    // Set tags for test for this test run:
    //
    // This writes all tags as one field:
    const params = {
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: 'tags',
        tags: this.tags.map(x => x.name + ':' + x.value).sort().join(',')
      }
    };
    await this.ddb.put(params).promise();

    await this.setTagsSearchStrings(this.tags);
    //
    // Update global list of tags:
    //
    const tagData = this.tags.map((x) => {
      const r = {
          Item: {
            testRunId: `__artilleryio:virtual:key`,
            kind: `tag#${x.name}:${x.value}`,
            name: x.name,
            valueType: typeof x.value,
          }
        };
      r.Item[`${(typeof x.value)}Value`] = x.value;
      return r;
    });
    debug(tagData);

    await this._writeInBatches(tagData);
    return this;
  }

  async setReport({aggregate}) {
    debug('testRunObject setReport');
    if (aggregate) {

      const aggregateJson = JSON.stringify(aggregate);
      debug('writing aggregate, item size:', aggregateJson.length);
      // Aggregate report:

      const objectPath = `${this.objectPrefix}/${this.id}/aggregate.json`;
      await this.objectStore.put(objectPath, aggregateJson);
      const params = {
        TableName: this.tableName,
        Item: {
          testRunId: this.id,
          kind: 'report#aggregate',
          data: `s3://${objectPath}`,
        }
      };
      await this.ddb.put(params).promise();
    }

    return this;
  }

  async setTextLog(lines) {
    debug('testRunObject setTextLog');

    let data;
    try {
      JSON.stringify(lines);
    } catch (stringifyErr) {
      throw new ConsoleOutputSerializeError('Failed trying to serialize console output');
    }

    let textLog = '';
    for(const args of lines) {
      textLog += util.format( ...Object.keys(args).map(k => args[k]))+'\n';
    }

    const objectPath = `${this.objectPrefix}/${this.id}/console-log.json`;
    await this.objectStore.put(objectPath, textLog);
    const params = {
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: 'consoleOutput#0',
        data: `s3://${objectPath}`,
      }
    };
    await this.ddb.put(params).promise();

    return this;
  }

  async setStatus(statusString, endTime) {
    debug('testRunObject setStatus:', statusString);

    const endTimeISOString = new Date(endTime).toISOString();

    await this.ddb.put({
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: 'lastStatus',
        statusName: statusString,
        createdTime: this.createdTime,
        endTime: endTimeISOString // TODO: endTime -> updatedTime
      }
    }).promise();

    // every time the status is updated, update it in the search facets too
    await this.setTagsSearchStrings(this.tags, {
      tagStatus: statusString,
      endTime: endTimeISOString
    });

    return this;
  }

  async setMetadata(metadata) {
    debug('testRunObject setMetadata');
    const items = [];
    // FIXME: this means metadata has to be a FLAT object, nothing
    // nested inside
    const metadata2 = JSON.parse(JSON.stringify(metadata));
    delete metadata2.tags;

    for (const [key, val] of Object.entries(metadata2)) {
      const r = {
        Item: {
          testRunId: this.id,
          kind: `metadata#${key}`,
          key: key,
          valueType: typeof val
        }
      };
      r.Item[`${typeof val}Value`] = val;
      items.push(r);
    }

    debug(items);

    await this._writeInBatches(items);

    this.createdTime = new Date(metadata.startedAt || Date.now()).toISOString();

    // Set createdTime attribute from metadata
    await this.ddb.put({
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: 'createdTime',
        createdTime: this.createdTime,
      }
    }).promise();

    return this;
  }

  async setTasks(taskArns) {
    debug('testRunObject setTasks');

    const items = taskArns.map((arn) => {
      return {
        Item: {
          testRunId: this.id,
          kind: `metadata#task#${arn}`,
          value: arn
        }
      }
    });

    await this._writeInBatches(items);
    return this;

  }

  // TODO: Refactor this + setReport()
  // ts is an epoch timestamp
  async recordIntermediateReport(report, ts) {
    debug('testRunObject recordIntermediateReport');

    await this.ddb.put({
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: `report#intermediate#${ts}`,
        data: JSON.stringify(report),
      }
    }).promise();

    return this;
  }

  // TODO: Add size checks
  async addNote(note) {
    debug('testRunObject addNote');

    // Add note to test:
    const ts = Date.now();
    await this.ddb.put({
      TableName: this.tableName,
      Item: {
        testRunId: this.id,
        kind: `note#${ts}`,
        noteId: uuidv4(),
        createdTime: new Date(ts).toISOString(),
        data: note // FIXME: XSS
      },
    }).promise();

    return this;
  }

  _makeBatchWriteRequest(tableName, items) {
    const batchWriteRequest = {
      RequestItems: {}
    };
    batchWriteRequest.RequestItems[tableName] = items.map(i => {
      return { PutRequest: i };
    });
    return batchWriteRequest;
  }

  _generateTagCombinations(tags) {
    const combinations = new Set();
    const slent = Math.pow(2, tags.length);
    let temp = [];

    for (var i = 0; i < slent; i++) {
      temp = [];
      for (var j = 0; j < tags.length; j++) {
        if (i & Math.pow(2, j)) {
          temp.push(tags[j]);
        }
      }

      if (temp.length > 0) {
        combinations.add(temp.sort().join(","));
      }
    }

    return [...combinations];
  }

  async _writeInBatches(items) {
    return writeInBatches(this.ddb, items, (batch) =>
      this._makeBatchWriteRequest(this.tableName, batch)
    );
  }
}

class StorageBackendNoop {
  constructor() {
    return this;
  }

  async init() {
    return this;
  }

  async setChecks() {
    return this;
  }

  async setEndedAt() {
    return this;
  }

  async setTags() {
    return this;
  }

  async bulkInsertTags() {
    return this;
  }

  async setTags() {
    return this;
  }

  async setReport() {
    return this;
  }

  async setTextLog() {
    return this;
  }

  async setStatus() {
    return this;
  }

  async setMetadata() {
    return this;
  }

  async setTasks() {
    return this;
  }

  async recordIntermediateReport() {
    return this;
  }

  async addNote() {
    return this;
  }
}

class TestRunObject {
  constructor(id, opts = {}) {
    this.id = id;
    if (opts && opts.type === 'aws:aurora-v1') {
      this.backend = new StorageBackendAuroraV1(id, opts);
    } else if (opts && opts.type == 'cloud') {
      this.backend = new StorageBackendCloud(id, opts);
    }else if (opts && opts.type == 'noop') {
      this.backend = new StorageBackendNoop(id, opts);
    } else {
      this.backend = new StorageBackendDynamo(id, opts);
    }

    this.opts = opts;
    return this;
  }

  async init(id, opts) {
    await this.backend.init(this.id, opts); // TODO: id is unnecessary
    return this;
  }

  static async get(id) {
  }

  static async list(attrs, limit, offset) {
    // supported attrs: id, tags
  }

  async setChecks(checks) {
    await this.backend.setChecks(checks);
    return this;
  }

  async setEndedAt(ts) {
    this.endedAt = ts;
    await this.backend.setEndedAt(ts);
    return this;
  }

  async bulkInsertTags(tags) {
    await this.backend.bulkInsertTags(tags);
    return this;
  }

  async setTags(tags) {
    this.tags = tags;
    await this.backend.setTags(tags);
    return this;
  }

  // TODO: This updates in one go, we want to do partial updates for intermediate
  async setReport(intermediateReports, aggregateReport) {
    this.report = {
      intermediate: intermediateReports,
      aggregate: aggregateReport
    };

    await this.backend.setReport(this.report);
    return this;
  }

  async setTextLog(lines) {
    this.textLog = lines;
    await this.backend.setTextLog(lines);
    return this;
  }

  async setStatus(statusString, ts) {
    if(typeof ts === 'undefined') {
      ts = Date.now();
    }
    await this.backend.setStatus(statusString, ts);
    return this;
  }

  async setMetadata(metadata) {
    await this.backend.setMetadata(metadata);
    return this;
  }

  async setTasks(taskArns) {
    await this.backend.setTasks(taskArns);
    return this;
  }

  async setTagsSearchStrings(tags, additionalAttributes) {
    await this.backend.setTagsSearchStrings(tags, additionalAttributes);
    return this;
  }

  async recordIntermediateReport(report, ts) {
    await this.backend.recordIntermediateReport(report, ts);
    return this;
  }

  async addNote(note, userId) {
    await this.backend.addNote(note, userId);
    return this;
  }
}

module.exports = {
  TestRunObject,
};