feat(schema-compiler): Support adding pre-aggregations in Rollup Designer for YAML-based models (#9500)

* prepare schema converter for yaml

* add yaml package

* fix: do not write all files while you need only one

* feat(schema-compiler): Support adding pre-aggregations in Rollup Designer for YAML-based models

* process and write only one required file

* add basic test for CubeSchemaConverter

* test for yaml pre-agg gen

* add note

* more tests

* rename test

* fix tests

Author: KSDaemon

packages/cubejs-schema-compiler/package.json

@@ -55,7 +55,8 @@
     "ramda": "^0.27.2",
     "syntax-error": "^1.3.0",
     "uuid": "^8.3.2",
-    "workerpool": "^9.2.0"
+    "workerpool": "^9.2.0",
+    "yaml": "^2.7.1"
   },
   "devDependencies": {
     "@clickhouse/client": "^1.7.0",

packages/cubejs-schema-compiler/src/compiler/converters/CubePreAggregationConverter.ts

@@ -1,9 +1,10 @@
 import { parse } from '@babel/parser';
 import * as t from '@babel/types';
+import YAML, { isMap, isScalar, Scalar, YAMLMap, YAMLSeq, Pair, parseDocument } from 'yaml';
 
 import { UserError } from '../UserError';
 
-import { AstByCubeName, CubeConverterInterface } from './CubeSchemaConverter';
+import { AstByCubeName, JsSet, CubeConverterInterface, YamlSet } from './CubeSchemaConverter';
 
 export type PreAggregationDefinition = {
   cubeName: string;
@@ -15,8 +16,20 @@ export class CubePreAggregationConverter implements CubeConverterInterface {
   public constructor(protected preAggregationDefinition: PreAggregationDefinition) {}
 
   public convert(astByCubeName: AstByCubeName): void {
-    const { cubeName, preAggregationName, code } = this.preAggregationDefinition;
-    const { cubeDefinition } = astByCubeName[cubeName];
+    const { cubeName } = this.preAggregationDefinition;
+
+    const cubeDefSet = astByCubeName[cubeName];
+
+    if ('ast' in cubeDefSet) {
+      this.convertJS(cubeDefSet);
+    } else {
+      this.convertYaml(cubeDefSet);
+    }
+  }
+
+  protected convertJS(cubeDefSet: JsSet) {
+    const { preAggregationName, code } = this.preAggregationDefinition;
+    const { cubeDefinition } = cubeDefSet;
 
     let preAggregationNode: t.ObjectExpression | null = null;
     const preAggregationAst = parse(`(${code})`);
@@ -64,4 +77,50 @@ export class CubePreAggregationConverter implements CubeConverterInterface {
       );
     }
   }
+
+  protected convertYaml(cubeDefSet: YamlSet) {
+    const { preAggregationName, code } = this.preAggregationDefinition;
+    const { cubeDefinition } = cubeDefSet;
+
+    const preAggDoc = YAML.parseDocument(code);
+    const preAggNode = preAggDoc.contents;
+
+    if (!preAggNode || !isMap(preAggNode)) {
+      throw new UserError('Pre-aggregation YAML must be a map/object');
+    }
+
+    const preAggsPair = cubeDefinition.items.find(
+      (pair: Pair) => isScalar(pair.key) && (pair.key.value === 'pre_aggregations' || pair.key.value === 'preAggregations')
+    );
+
+    if (preAggsPair) {
+      const seq = preAggsPair.value;
+      if (!YAML.isSeq(seq)) {
+        throw new UserError('\'pre_aggregations\' must be a sequence');
+      }
+
+      const exists = seq.items.some(item => {
+        if (isMap(item)) {
+          const namePair = item.items.find(
+            (pair: Pair) => isScalar(pair.key) && pair.key.value === 'name'
+          );
+          return namePair && isScalar(namePair.value) && namePair.value.value === preAggregationName;
+        }
+        return false;
+      });
+
+      if (exists) {
+        throw new UserError(`Pre-aggregation '${preAggregationName}' is already defined`);
+      }
+
+      seq.items.push(preAggNode);
+    } else {
+      const newSeq = new YAMLSeq();
+      newSeq.items.push(preAggNode);
+
+      cubeDefinition.items.push(
+        new Pair(new Scalar('pre_aggregations'), newSeq)
+      );
+    }
+  }
 }

packages/cubejs-schema-compiler/src/compiler/converters/CubeSchemaConverter.ts

@@ -2,14 +2,23 @@ import generator from '@babel/generator';
 import { parse } from '@babel/parser';
 import traverse from '@babel/traverse';
 import * as t from '@babel/types';
+import YAML, { isMap, isSeq } from 'yaml';
 
-export type AstSet = {
+export type JsSet = {
   fileName: string;
   ast: t.File;
   cubeDefinition: t.ObjectExpression;
 };
 
-export type AstByCubeName = Record<string, AstSet>;
+export type YamlSet = {
+  fileName: string;
+  yaml: YAML.Document;
+  cubeDefinition: YAML.YAMLMap;
+};
+
+const JINJA_SYNTAX = /{%|%}|{{|}}/ig;
+
+export type AstByCubeName = Record<string, (JsSet | YamlSet)>;
 
 export interface CubeConverterInterface {
   convert(astByCubeName: AstByCubeName): void;
@@ -27,47 +36,112 @@ export class CubeSchemaConverter {
 
   public constructor(protected fileRepository: any, protected converters: CubeConverterInterface[]) {}
 
-  protected async prepare(): Promise<void> {
+  /**
+   * Parse Schema files from the repository and create a mapping of cube names to schema files.
+   * If optional `cubeName` parameter is passed - only file with asked cube is parsed and stored.
+   * @param cubeName
+   * @protected
+   */
+  protected async prepare(cubeName?: string): Promise<void> {
     this.dataSchemaFiles = await this.fileRepository.dataSchemaFiles();
 
     this.dataSchemaFiles.forEach((schemaFile) => {
-      const ast = this.parse(schemaFile);
-
-      traverse(ast, {
-        CallExpression: (path) => {
-          if (t.isIdentifier(path.node.callee)) {
-            const args = path.get('arguments');
-
-            if (path.node.callee.name === 'cube') {
-              if (args?.[args.length - 1]) {
-                let cubeName: string | undefined;
-
-                if (args[0].node.type === 'StringLiteral' && args[0].node.value) {
-                  cubeName = args[0].node.value;
-                } else if (args[0].node.type === 'TemplateLiteral' && args[0].node.quasis?.[0]?.value.cooked) {
-                  cubeName = args[0].node.quasis?.[0]?.value.cooked;
-                }
-
-                if (cubeName == null) {
-                  throw new Error(`Error parsing ${schemaFile.fileName}`);
-                }
-
-                if (t.isObjectExpression(args[1]?.node) && ast != null) {
-                  this.parsedFiles[cubeName] = {
-                    fileName: schemaFile.fileName,
-                    ast,
-                    cubeDefinition: args[1].node,
-                  };
-                }
+      if (schemaFile.fileName.endsWith('.js')) {
+        this.transformJS(schemaFile, cubeName);
+      } else if ((schemaFile.fileName.endsWith('.yml') || schemaFile.fileName.endsWith('.yaml')) && !schemaFile.content.match(JINJA_SYNTAX)) {
+        // Jinja-templated data models are not supported in Rollup Designer yet, so we're ignoring them,
+        // and if user has chosen the cube from such file - it won't be found during generation.
+        this.transformYaml(schemaFile, cubeName);
+      }
+    });
+  }
+
+  protected transformYaml(schemaFile: SchemaFile, filterCubeName?: string) {
+    if (!schemaFile.content.trim()) {
+      return;
+    }
+
+    const yamlDoc = YAML.parseDocument(schemaFile.content);
+    if (!yamlDoc?.contents) {
+      return;
+    }
+
+    const root = yamlDoc.contents;
+
+    if (!isMap(root)) {
+      return;
+    }
+
+    const cubesPair = root.items.find((item) => {
+      const key = item.key as YAML.Scalar;
+      return key?.value === 'cubes';
+    });
+
+    if (!cubesPair || !isSeq(cubesPair.value)) {
+      return;
+    }
+
+    for (const cubeNode of cubesPair.value.items) {
+      if (isMap(cubeNode)) {
+        const cubeNamePair = cubeNode.items.find((item) => {
+          const key = item.key as YAML.Scalar;
+          return key?.value === 'name';
+        });
+
+        const cubeName = (cubeNamePair?.value as YAML.Scalar).value;
+
+        if (cubeName && typeof cubeName === 'string' && (!filterCubeName || cubeName === filterCubeName)) {
+          this.parsedFiles[cubeName] = {
+            fileName: schemaFile.fileName,
+            yaml: yamlDoc,
+            cubeDefinition: cubeNode,
+          };
+
+          if (cubeName === filterCubeName) {
+            return;
+          }
+        }
+      }
+    }
+  }
+
+  protected transformJS(schemaFile: SchemaFile, filterCubeName?: string) {
+    const ast = this.parseJS(schemaFile);
+
+    traverse(ast, {
+      CallExpression: (path) => {
+        if (t.isIdentifier(path.node.callee)) {
+          const args = path.get('arguments');
+
+          if (path.node.callee.name === 'cube') {
+            if (args?.[args.length - 1]) {
+              let cubeName: string | undefined;
+
+              if (args[0].node.type === 'StringLiteral' && args[0].node.value) {
+                cubeName = args[0].node.value;
+              } else if (args[0].node.type === 'TemplateLiteral' && args[0].node.quasis?.[0]?.value.cooked) {
+                cubeName = args[0].node.quasis?.[0]?.value.cooked;
+              }
+
+              if (cubeName == null) {
+                throw new Error(`Error parsing ${schemaFile.fileName}`);
+              }
+
+              if (t.isObjectExpression(args[1]?.node) && ast != null && (!filterCubeName || cubeName === filterCubeName)) {
+                this.parsedFiles[cubeName] = {
+                  fileName: schemaFile.fileName,
+                  ast,
+                  cubeDefinition: args[1].node,
+                };
               }
             }
           }
-        },
-      });
+        }
+      },
     });
   }
 
-  protected parse(file: SchemaFile) {
+  protected parseJS(file: SchemaFile) {
     try {
       return parse(file.content, {
         sourceFilename: file.fileName,
@@ -86,19 +160,25 @@ export class CubeSchemaConverter {
     }
   }
 
-  public async generate() {
-    await this.prepare();
+  public async generate(cubeName?: string) {
+    await this.prepare(cubeName);
 
     this.converters.forEach((converter) => {
       converter.convert(this.parsedFiles);
     });
   }
 
   public getSourceFiles() {
-    return Object.entries(this.parsedFiles).map(([cubeName, file]) => ({
-      cubeName,
-      fileName: file.fileName,
-      source: generator(file.ast, {}).code,
-    }));
+    return Object.entries(this.parsedFiles).map(([cubeName, file]) => {
+      const source = 'ast' in file
+        ? generator(file.ast, {}).code
+        : String(file.yaml);
+
+      return {
+        cubeName,
+        fileName: file.fileName,
+        source,
+      };
+    });
   }
 }