@getanthill/datastore/src/sdk/aggregator/Aggregator.ts

Event-Sourced Datastore

import type Datastore from '../Datastore';

import crypto from 'crypto';

import Ajv from 'ajv';
import addFormats from 'ajv-formats';

import get from 'lodash/get';
import set from 'lodash/set';
import unset from 'lodash/unset';
import mergeWith from 'lodash/mergeWith';
import cloneDeep from 'lodash/cloneDeep';

import * as jsonpatch from 'fast-json-patch';

import schemas from '../schemas.json';

import lodash from 'lodash';

import * as utils from '../../utils';
import { MultiQuery } from '../utils';

const validator = new Ajv({
  useDefaults: false,
  coerceTypes: true,
  strict: false,
});

// @ts-ignore
addFormats(validator);

/**
 * Type of steps
 * - fetch
 * - map
 * - unset
 * - validate
 * - each
 * - filter
 * - op
 *
 * - parallel
 * - request
 * - json_patch
 */
export const STEP_TYPE_FETCH = 'fetch';
export const STEP_TYPE_JSON_PATCH = 'json_patch';
export const STEP_TYPE_MAP = 'map';
export const STEP_TYPE_UNSET = 'unset';
export const STEP_TYPE_VALIDATE = 'validate';
export const STEP_TYPE_EACH = 'each';
export const STEP_TYPE_IF = 'if';
export const STEP_TYPE_FILTER = 'filter';
export const STEP_TYPE_OP = 'op';
export const STEP_TYPE_PERSIST = 'persist';
export const STEP_TYPE_FROM = 'from';
export const DEFAULT_STEP_TYPES = Object.freeze([
  STEP_TYPE_FETCH,
  STEP_TYPE_JSON_PATCH,
  STEP_TYPE_MAP,
  STEP_TYPE_UNSET,
  STEP_TYPE_VALIDATE,
  STEP_TYPE_EACH,
  STEP_TYPE_IF,
  STEP_TYPE_FILTER,
  STEP_TYPE_OP,
  STEP_TYPE_PERSIST,
  STEP_TYPE_FROM,
]);

export interface StepDefinition {
  handler: Function;
  schema?: object;

  name?: string;
  description?: string;
}

export interface StepMapItem {
  from: string;
  to: string;
  default: any;
  must_hash?: boolean;
  json_stringify?: boolean;
  relative_date_in_seconds?: number;
}

export interface StepFetch {
  type: string;
  name?: string;
  description?: string;

  destination?: string;

  datastore?: string;
  model: string;
  source?: MultiQuery['source'];

  query?: object;
  headers?: any;
  map?: StepMapItem[];

  timetravel?: string;
  correlation_field?: string;

  page?: number;
  page_size?: number;

  as_entity?: boolean;
  must_decrypt?: boolean;
  default?: any;
}

export interface StepPersist {
  type: string;
  name?: string;
  description?: string;

  destination: string;

  datastore?: string;
  model: string;

  correlation_field?: string;

  payload: any;
  headers?: any;
  map?: StepMapItem[];
  imperative_version_next?: string;
}

export interface StepJsonPatch {
  type: string;
  name?: string;
  description?: string;

  patch: jsonpatch.Operation[];
}

export interface StepMap {
  type: string;
  name?: string;
  description?: string;

  map: StepMapItem[];
}

/**
 * @deprecated in favor of JSON Patch
 */
export interface StepUnset {
  type: string;
  name?: string;
  description?: string;

  path: string;
}

export interface StepValidate {
  type: string;
  name?: string;
  description?: string;

  schema: object;

  path?: string;
  must_throw?: boolean;
  destination?: string;
}

export interface StepEach {
  type: string;
  name?: string;
  description?: string;

  path: string;
  pipeline: Step[];

  destination?: string;
}

export interface StepIf {
  type: string;
  name?: string;
  description?: string;

  path: string;
  schema: object;
  pipeline: Step[];

  destination?: string;
  repeat_while_true?: boolean;
  max_iteration_count?: number;
}

export interface StepFilter {
  type: string;
  name?: string;
  description?: string;

  path: string;
  schema: object;

  destination?: string;
  map?: StepMapItem[];
  as_entity?: boolean;
  default?: any;
}

export interface StepOpArg {
  path: string;
  default: number;
  func: string;
}

export interface StepOp {
  type: string;
  name?: string;
  description?: string;

  func: string;
  args: StepOpArg[];

  path?: string;
  destination?: string;

  default?: any;
  args_as_array?: boolean;
}

export interface StepFrom {
  type: string;
  name?: string;
  description?: string;

  path: string;
}

export type Step =
  | StepFetch
  | StepJsonPatch
  | StepMap
  | StepUnset
  | StepValidate
  | StepEach
  | StepIf
  | StepFilter
  | StepOp
  | StepFrom;

export type Aggregation = { [key: string]: any };

export default class Aggregator {
  datastores: Map<string, Datastore>;
  steps: Map<string, StepDefinition> = new Map();
  logs: { level: string; msg: string; ts: number; context?: any }[] = [];

  static pipelineValidator = new Ajv({
    schemas,
    useDefaults: false,
    coerceTypes: false,
    strict: true,
  });

  static ERROR_INVALID_PIPELINE_DEFINITION = new Error(
    'Invalid pipeline definition',
  );
  static ERROR_INVALID_STEP_TYPE = new Error('Invalid step type');
  static ERROR_ENTITY_NOT_FOUND = new Error('Entity not found');
  static ERROR_DESTINATION_UNDEFINED = new Error('Destination must be defined');
  static ERROR_CONFLICT_STEP_TYPE = new Error('Conflict on step type');
  static ERROR_CONTRACT_ERROR_STEP_TYPE = new Error(
    'Contract error on step type definition',
  );
  static ERROR_VALIDATE_STEP_FAILED = new Error('Validate step failed');
  static ERROR_ITERATE_STEP_IS_NOT_ARRAY = new Error(
    'Iterate step not on array',
  );
  static ERROR_OP_STEP_INVALID_VALUE = new Error('Op step value is invalid');
  static ERRORS = [
    Aggregator.ERROR_INVALID_STEP_TYPE,
    Aggregator.ERROR_ENTITY_NOT_FOUND,
    Aggregator.ERROR_DESTINATION_UNDEFINED,
    Aggregator.ERROR_CONFLICT_STEP_TYPE,
    Aggregator.ERROR_CONTRACT_ERROR_STEP_TYPE,
    Aggregator.ERROR_VALIDATE_STEP_FAILED,
    Aggregator.ERROR_ITERATE_STEP_IS_NOT_ARRAY,
    Aggregator.ERROR_OP_STEP_INVALID_VALUE,
  ];

  config: { max_retry: number };

  metrics: {
    started_at_in_ms?: number;
    ended_at_in_ms?: number;
    elapsed_time_in_ms?: number;
    iteration?: number;
  };

  private _pipelineValidator = Aggregator.pipelineValidator;

  constructor(
    datastores: Map<string, Datastore>,
    config: { max_retry: number } = { max_retry: 0 },
  ) {
    this.datastores = datastores;
    this.config = config;
    this.metrics = {};
  }

  static _customizer(objValue: any, srcValue: any) {
    if (Array.isArray(objValue)) {
      return objValue.concat(srcValue);
    }
  }

  get pipelineValidator() {
    return this._pipelineValidator;
  }

  log(level: string, msg: string, context?: any): void {
    this.logs.push({ ts: Date.now(), level, msg, context });
  }

  addStepType(stepType: string, stepDefinition: StepDefinition): this {
    if (this.steps.has(stepType)) {
      throw Aggregator.ERROR_CONFLICT_STEP_TYPE;
    }

    this.steps.set(stepType, stepDefinition);

    this.updateValidator();

    return this;
  }

  removeStepType(stepType: string): this {
    this.steps.delete(stepType);

    this.updateValidator();

    return this;
  }

  updateValidator() {
    const _schemas = cloneDeep(schemas);
    for (const key of this.steps.keys()) {
      _schemas[0].items?.oneOf?.push({
        $ref: `/schemas/datastore/aggregator/step/${key}`,
      });
      _schemas.push({
        $id: `/schemas/datastore/aggregator/step/${key}`,

        type: 'object',
        required: ['type'],
        // @ts-ignore
        properties: {
          type: {
            $ref: '/schemas/datastore/aggregator/components/type',
            enum: [key],
          },
        },
      });
    }

    this._pipelineValidator = new Ajv({
      schemas: _schemas,
      useDefaults: false,
      coerceTypes: false,
      strict: true,
    });
  }

  mergeData(
    step: StepFetch | StepFilter,
    data: Aggregation = {},
    results: object[] = step.default || [],
    destination = step.destination,
  ) {
    if (!destination) {
      throw Aggregator.ERROR_DESTINATION_UNDEFINED;
    }

    const res = step.as_entity === true ? results[0] || step.default : results;

    if (step.as_entity === true && res === undefined) {
      throw Aggregator.ERROR_ENTITY_NOT_FOUND;
    }

    return mergeWith(
      {},
      data,
      step.as_entity === true && destination === '.'
        ? res
        : set({}, destination, res),
      Aggregator._customizer,
    );
  }

  static ok(condition: boolean, message: string) {
    if (condition !== true) {
      const err = new Error(message);
      throw err;
    }
  }

  static hash(value: any): string {
    const hash = crypto.createHash('sha512');

    hash.update(value);

    return hash.digest('hex');
  }

  applyPatch(data: Aggregation, patch: jsonpatch.Operation[]) {
    const { newDocument } = jsonpatch.applyPatch(data, patch, true);

    return newDocument;
  }

  applyMap(
    source: object = {},
    map: StepMapItem[] = [],
    data: Aggregation | null = null,
  ): object {
    map.forEach(
      ({
        from,
        to,
        default: defaultValue,
        must_hash: mustHash,
        json_stringify: jsonStringify,
        relative_date_in_seconds: relativeDateInSeconds,
      }) => {
        let value =
          from === '.' ? (data ?? defaultValue) : get(data, from, defaultValue);

        if (relativeDateInSeconds) {
          const start = value ? new Date(value).getTime() : Date.now();
          value = new Date(start + relativeDateInSeconds * 1000);
        }

        if (jsonStringify === true) {
          value = JSON.stringify(value);
        }

        if (mustHash === true) {
          value = Aggregator.hash(value);
        }

        if (to === '.') {
          source = value;
        } else {
          set(source, to, value);
        }
      },
    );

    return source;
  }

  async fetch(step: StepFetch, data: Aggregation = {}): Promise<object[]> {
    const datastore = this.getDatastore(step.datastore);

    Aggregator.ok(typeof datastore.walk === 'function', 'Invalid datastore');

    let results: object[] = [];

    const query = this.applyMap(step.query, step.map, data);

    this.log('debug', 'Fetch query', {
      model: step.model,
      source: step.source,
      headers: step.headers,
      query,
    });

    if ('page' in step || 'page_size' in step) {
      if (step.source === 'events') {
        const { data } = await datastore.allEvents(
          step.model,
          query,
          step.page,
          step.page_size,
          step.headers,
        );

        results = data;
      } else {
        const { data } = await datastore.find(
          step.model,
          query,
          step.page,
          step.page_size,
          step.headers,
        );

        results = data;
      }
    } else {
      await datastore.walk(
        step.model,
        query,
        (entity: any) => {
          results.push(entity);
        },
        10,
        step.source,
        step.headers,
      );
    }

    const timetravelDate = get(data, step.timetravel ?? '');
    Aggregator.ok(
      step.timetravel === undefined ||
        (!!timetravelDate && !!step.correlation_field),
      'Invalid timetravel condition',
    );
    if (!!timetravelDate && !!step.correlation_field) {
      const res = await Promise.all(
        results.map((r: { [key: string]: any }) =>
          datastore.at(step.model, r[step.correlation_field!], timetravelDate),
        ),
      );

      results = res.map((r) => r.data).filter((v) => v);
    }

    if (step.must_decrypt === true) {
      const res = await datastore.decrypt(step.model, results);

      results = res.data;
    }

    return results;
  }

  getDatastore(datastore: string | undefined): Datastore {
    if (datastore) {
      const _datastore = this.datastores.get(datastore);

      if (_datastore) {
        return _datastore;
      }
    }

    const [_datastore] = this.datastores.values();

    return _datastore;
  }

  async persist(step: StepPersist, data: Aggregation = {}): Promise<any> {
    const datastore = this.getDatastore(step.datastore);

    Aggregator.ok(typeof datastore.walk === 'function', 'Invalid datastore');

    const payload: any = this.applyMap(step.payload, step.map, data);
    const headers = { ...(step.headers || {}) };

    if ('imperative_version_next' in step) {
      headers.version = (get(data, step.imperative_version_next!) ?? -1) + 1;
    }

    const correlationField = step.correlation_field ?? null;
    const correlationId =
      correlationField !== null ? payload[correlationField] : null;

    Aggregator.ok(
      !!correlationId || correlationId === null,
      'Correlation ID must be null or exist',
    );

    let result;
    if (correlationId === null) {
      if (payload.created_at) {
        headers['created-at'] = payload.created_at;

        delete payload.created_at;
        delete payload.updated_at;
      }

      const res = await datastore.create(step.model, payload, headers);

      result = res.data;
    } else {
      if ('updated_at' in payload || 'created_at' in payload) {
        headers['created-at'] = payload.updated_at || payload.created_at;

        delete payload.created_at;
        delete payload.updated_at;
      }

      const res = await datastore.update(
        step.model,
        correlationId,
        payload,
        headers,
      );

      result = res.data;
    }

    return result;
  }

  async runStepFetch(
    step: StepFetch,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const results = await this.fetch(step, data);

    return this.mergeData(step, data, results, step.destination ?? step.model);
  }

  async runStepPersist(
    step: StepPersist,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const res = await this.persist(step, data);

    const destination = step.destination || 'persist';

    set(data, destination, res);

    return data;
  }

  async runStepJsonPatch(
    step: StepJsonPatch,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    return this.applyPatch(data, step.patch);
  }

  async runStepMap(
    step: StepMap,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    return this.applyMap(data, step.map, data);
  }

  async runStepUnset(
    step: StepUnset,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    unset(data, step.path);

    return data;
  }

  async runStepValidate(
    step: StepValidate,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const value = step.path ? get(data, step.path) : data;
    const isValid = validator.validate(step.schema, value);

    this.log('debug', 'Validate', { step, value, is_valid: isValid });

    if (isValid === false && step.must_throw === true) {
      const err = new Error(Aggregator.ERROR_VALIDATE_STEP_FAILED.message);

      // @ts-ignore
      err.step = step;
      // @ts-ignore
      err.schema = step.schema;
      // @ts-ignore
      err.value = value;
      // @ts-ignore
      err.errors = validator.errors;

      throw err;
    }

    const destination = step.destination ?? 'validation';
    const res = { is_valid: isValid, errors: validator.errors };

    set(data, destination, res);

    return data;
  }

  async runStepEach(
    step: StepEach,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const items = get(data, step.path, null);

    if (!Array.isArray(items)) {
      throw Aggregator.ERROR_ITERATE_STEP_IS_NOT_ARRAY;
    }

    const aggregates = await Promise.all(
      items.map((item) => this.aggregate(step.pipeline, { ...item, _: data })),
    );

    const destination = step.destination ?? step.path;

    set(
      data,
      destination,
      aggregates.map((agg: any) => {
        delete agg._;

        return agg;
      }),
    );

    return data;
  }

  async runStepIf(
    step: StepIf,
    data: Aggregation = {},
    i = 0,
  ): Promise<Aggregation> {
    const maxIterationCount: number = step.max_iteration_count ?? 100;
    if (i >= maxIterationCount) {
      return data;
    }

    const value = step.path ? get(data, step.path) : data;
    const isValid = validator.validate(step.schema, value);

    if (isValid === false) {
      return data;
    }

    const _data = await this.aggregate(step.pipeline, data);

    if (step.repeat_while_true === true) {
      return this.runStepIf(step, _data, i + 1);
    }

    return _data;
  }

  async runStepFilter(
    step: StepFilter,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const items = get(data, step.path, null);

    if (!Array.isArray(items)) {
      throw Aggregator.ERROR_ITERATE_STEP_IS_NOT_ARRAY;
    }

    const schema = this.applyMap(step.schema, step.map, data);

    const results = items.filter((item) =>
      validator.validate(schema, item),
    ) as any[];

    if (!step.destination) {
      unset(data, step.path);
    }

    return this.mergeData(step, data, results, step.destination ?? step.path);
  }

  /**
   * @alpha
   */
  async runStepOp(
    step: StepOp,
    data: Aggregation | null = null,
  ): Promise<Aggregation> {
    const value = step.path
      ? get(data, step.path, step.default)
      : (step.default ?? data);

    let args =
      'args' in step
        ? step.args.map((arg: StepOpArg) => {
            if (typeof arg !== 'object') {
              return arg;
            }

            if (arg.func) {
              // @ts-ignore
              return lodash[arg.func];
            }

            if (arg.path === '.') {
              return value;
            }

            return get(value, arg.path, arg.default);
          })
        : [value];

    if (step.args_as_array === true) {
      args = [args];
    }

    let res: any = null;

    if (step.func === 'length') {
      res = args[0].length;
    } else if (step.func === 'date') {
      res = utils.getDate(args[0]).toISOString();
    } else if (step.func === 'timestamp') {
      res = utils.getDate(args[0]).getTime();
    } else {
      /* @ts-ignore */
      res = lodash[step.func].apply(null, args);
    }

    const destination = step.destination ?? step.func;

    data = data === null ? {} : data;
    if (destination === '.') {
      return res;
    }

    set(data, destination, res);

    return data;
  }

  async runStepFrom(
    step: StepFrom,
    data: Aggregation = {},
  ): Promise<Aggregation> {
    const steps = get(data, step.path, []);

    return this.aggregate(steps, data);
  }

  async runStep(step: Step, data: Aggregation = {}): Promise<Aggregation> {
    utils.mapValuesDeep(step, (v: unknown) =>
      typeof v === 'string' && v.startsWith('{') && v.endsWith('}')
        ? cloneDeep(get(data, v.slice(1, -1), v))
        : v,
    );

    const stepDefinition = this.steps.get(step.type) as StepDefinition;

    if (stepDefinition !== undefined) {
      const stepHandler = stepDefinition.handler;
      Aggregator.ok(
        typeof stepHandler === 'function',
        'Invalid step handler defined',
      );

      if (
        !!stepDefinition.schema &&
        validator.validate(stepDefinition.schema, step) !== true
      ) {
        throw Aggregator.ERROR_CONTRACT_ERROR_STEP_TYPE;
      }

      return stepHandler.call(this, step, data);
    }

    if (step.type === STEP_TYPE_FETCH) {
      return this.runStepFetch(step as StepFetch, data);
    }

    if (step.type === STEP_TYPE_JSON_PATCH) {
      return this.runStepJsonPatch(step as StepJsonPatch, data);
    }

    if (step.type === STEP_TYPE_MAP) {
      return this.runStepMap(step as StepMap, data);
    }

    if (step.type === STEP_TYPE_UNSET) {
      return this.runStepUnset(step as StepUnset, data);
    }

    if (step.type === STEP_TYPE_VALIDATE) {
      return this.runStepValidate(step as StepValidate, data);
    }

    if (step.type === STEP_TYPE_EACH) {
      return this.runStepEach(step as StepEach, data);
    }

    if (step.type === STEP_TYPE_IF) {
      return this.runStepIf(step as StepIf, data);
    }

    if (step.type === STEP_TYPE_FILTER) {
      return this.runStepFilter(step as StepFilter, data);
    }

    if (step.type === STEP_TYPE_OP) {
      return this.runStepOp(step as StepOp, data);
    }

    if (step.type === STEP_TYPE_PERSIST) {
      return this.runStepPersist(step as StepPersist, data);
    }

    if (step.type === STEP_TYPE_FROM) {
      return this.runStepFrom(step as StepFrom, data);
    }

    throw Aggregator.ERROR_INVALID_STEP_TYPE;
  }

  validate(pipeline: Step[]) {
    const isValid = this.pipelineValidator.validate(
      '/schemas/datastore/aggregator/pipeline',
      pipeline,
    );

    if (isValid === false) {
      const err = Aggregator.ERROR_INVALID_PIPELINE_DEFINITION;

      // @ts-ignore
      err.details = this.pipelineValidator.errors;

      throw err;
    }
  }

  async aggregate(
    pipeline: Step[],
    initialData: Aggregation = {},
    iteration = 0,
  ): Promise<Aggregation> {
    this.validate(pipeline);

    this.logs = [];
    const _pipeline = cloneDeep(pipeline);
    let data = cloneDeep(initialData);

    this.metrics.started_at_in_ms = this.metrics.started_at_in_ms ?? Date.now();
    this.metrics.iteration = iteration;

    try {
      for (const step of _pipeline) {
        this.log('debug', 'Before step', { step, data });
        data = await this.runStep(step, data);
        this.log('debug', 'After step', { step, data });
      }
    } catch (err) {
      if (iteration > (this.config.max_retry || 0)) {
        this.log('error', 'Max retry reached', {
          err,
          iteration,
          max_retry: this.config.max_retry || 0,
        });

        throw err;
      }

      this.log('error', 'Retrying the aggregation', {
        err,
        iteration,
        max_retry: this.config.max_retry || 0,
      });

      return this.aggregate(pipeline, initialData, iteration + 1);
    }

    this.metrics.ended_at_in_ms = Date.now();
    this.metrics.elapsed_time_in_ms =
      this.metrics.ended_at_in_ms - this.metrics.started_at_in_ms;

    this.log('debug', 'Aggregation ended successfully', {
      metrics: this.metrics,
    });

    return data;
  }
}