feat(schema-compiler): composite key generation (#7929)

cube-js · Mar 15, 2024 · 5a27200 · 5a27200
1 parent 7f09fe8
commit 5a27200
Show file tree

Hide file tree

Showing 3 changed files with 202 additions and 36 deletions.
diff --git a/packages/cubejs-schema-compiler/src/scaffolding/ScaffoldingSchema.ts b/packages/cubejs-schema-compiler/src/scaffolding/ScaffoldingSchema.ts
@@ -53,6 +53,7 @@ export type CubeDescriptor = {
   schema: string;
   members: CubeDescriptorMember[];
   joins: Join[];
+  shouldGeneratePrimaryKey?: boolean;
 };
 
 export type TableSchema = {
@@ -64,6 +65,7 @@ export type TableSchema = {
   dimensions: Dimension[];
   drillMembers?: Dimension[];
   joins: Join[];
+  compositePrimaryKey?: any[];
 };
 
 const MEASURE_DICTIONARY = [

diff --git a/packages/cubejs-schema-compiler/src/scaffolding/formatters/BaseSchemaFormatter.ts b/packages/cubejs-schema-compiler/src/scaffolding/formatters/BaseSchemaFormatter.ts
@@ -1,7 +1,8 @@
-import inflection from 'inflection';
+import inflection, { titleize } from 'inflection';
 import { CubeMembers, SchemaContext } from '../ScaffoldingTemplate';
 import {
   CubeDescriptor,
+  CubeDescriptorMember,
   DatabaseSchema,
   MemberType,
   ScaffoldingSchema,
@@ -31,9 +32,9 @@ export abstract class BaseSchemaFormatter {
   protected readonly scaffoldingSchema: ScaffoldingSchema;
 
   public constructor(
-      protected readonly dbSchema: DatabaseSchema,
-      protected readonly driver: any,
-      protected readonly options: SchemaFormatterOptions
+    protected readonly dbSchema: DatabaseSchema,
+    protected readonly driver: any,
+    protected readonly options: SchemaFormatterOptions
   ) {
     this.scaffoldingSchema = new ScaffoldingSchema(dbSchema, this.options);
   }
@@ -42,7 +43,9 @@ export abstract class BaseSchemaFormatter {
 
   protected abstract cubeReference(cube: string): string;
 
-  protected abstract renderFile(fileDescriptor: Record<string, unknown>): string;
+  protected abstract renderFile(
+    fileDescriptor: Record<string, unknown>
+  ): string;
 
   public generateFilesByTableNames(
     tableNames: TableName[],
@@ -54,18 +57,24 @@ export abstract class BaseSchemaFormatter {
 
     return schemaForTables.map((tableSchema) => ({
       fileName: `${tableSchema.cube}.${this.fileExtension()}`,
-      content: this.renderFile(this.schemaDescriptorForTable(tableSchema, schemaContext)),
+      content: this.renderFile(
+        this.schemaDescriptorForTable(tableSchema, schemaContext)
+      ),
     }));
   }
 
   public generateFilesByCubeDescriptors(
     cubeDescriptors: CubeDescriptor[],
     schemaContext: SchemaContext = {}
   ): SchemaFile[] {
-    return this.schemaForTablesByCubeDescriptors(cubeDescriptors).map((tableSchema) => ({
-      fileName: `${tableSchema.cube}.${this.fileExtension()}`,
-      content: this.renderFile(this.schemaDescriptorForTable(tableSchema, schemaContext)),
-    }));
+    return this.schemaForTablesByCubeDescriptors(cubeDescriptors).map(
+      (tableSchema) => ({
+        fileName: `${tableSchema.cube}.${this.fileExtension()}`,
+        content: this.renderFile(
+          this.schemaDescriptorForTable(tableSchema, schemaContext)
+        ),
+      })
+    );
   }
 
   protected sqlForMember(m) {
@@ -77,7 +86,8 @@ export abstract class BaseSchemaFormatter {
   }
 
   protected memberTitle(m) {
-    return inflection.titleize(inflection.underscore(this.memberName(m))) !== m.title
+    return inflection.titleize(inflection.underscore(this.memberName(m))) !==
+      m.title
       ? m.title
       : undefined;
   }
@@ -103,11 +113,16 @@ export abstract class BaseSchemaFormatter {
     return !!name.match(/^[a-z0-9_]+$/i);
   }
 
-  public schemaDescriptorForTable(tableSchema: TableSchema, schemaContext: SchemaContext = {}) {
+  public schemaDescriptorForTable(
+    tableSchema: TableSchema,
+    schemaContext: SchemaContext = {}
+  ) {
     let table = `${
-      tableSchema.schema?.length ? `${this.escapeName(tableSchema.schema)}.` : ''
+      tableSchema.schema?.length
+        ? `${this.escapeName(tableSchema.schema)}.`
+        : ''
     }${this.escapeName(tableSchema.table)}`;
-    
+
     if (this.options.catalog) {
       table = `${this.escapeName(this.options.catalog)}.${table}`;
     }
@@ -116,7 +131,9 @@ export abstract class BaseSchemaFormatter {
 
     let dataSourceProp = {};
     if (dataSource) {
-      dataSourceProp = this.options.snakeCase ? { data_source: dataSource } : { dataSource };
+      dataSourceProp = this.options.snakeCase
+        ? { data_source: dataSource }
+        : { dataSource };
     }
 
     const sqlOption = this.options.snakeCase
@@ -126,7 +143,24 @@ export abstract class BaseSchemaFormatter {
       : {
         sql: `SELECT * FROM ${table}`,
       };
-
+
+    let compositePrimaryKey: any;
+
+    if (tableSchema.compositePrimaryKey) {
+      const { length } = tableSchema.compositePrimaryKey;
+      const name = tableSchema.dimensions.find(d => d.name === 'id') ? 'generated_composite_primary_key' : 'id';
+      compositePrimaryKey = {
+        [name]: {
+          name,
+          title: titleize(name),
+          sql: this.driver.concatStringsSql(tableSchema.compositePrimaryKey
+            .flatMap(
+              (it, index) => [`${this.cubeReference('CUBE')}.${this.escapeName(it.name)}`, index !== length - 1 ? '\'-\'' : null].filter(Boolean)
+            )),
+          [this.options.snakeCase ? 'primary_key' : 'primaryKey']: true
+        } };
+    }
+
     return {
       cube: tableSchema.cube,
       ...sqlOption,
@@ -137,25 +171,30 @@ export abstract class BaseSchemaFormatter {
           [j.cubeToJoin]: {
             sql: `${this.cubeReference('CUBE')}.${this.escapeName(
               j.thisTableColumn
-            )} = ${this.cubeReference(j.cubeToJoin)}.${this.escapeName(j.columnToJoin)}`,
+            )} = ${this.cubeReference(j.cubeToJoin)}.${this.escapeName(
+              j.columnToJoin
+            )}`,
             relationship: this.options.snakeCase
               ? JOIN_RELATIONSHIP_MAP[j.relationship]
               : j.relationship,
           },
         }))
         .reduce((a, b) => ({ ...a, ...b }), {}),
-      dimensions: tableSchema.dimensions.sort((a) => (a.isPrimaryKey ? -1 : 0))
-        .map((m) => ({
-          [this.memberName(m)]: {
-            sql: this.sqlForMember(m),
-            type: m.type ?? m.types[0],
-            title: this.memberTitle(m),
-            [this.options.snakeCase ? 'primary_key' : 'primaryKey']: m.isPrimaryKey
-              ? true
-              : undefined,
-          },
-        }))
-        .reduce((a, b) => ({ ...a, ...b }), {}),
+      dimensions: {
+        ...compositePrimaryKey,
+        ...tableSchema.dimensions
+          .sort((a) => (a.isPrimaryKey ? -1 : 0))
+          .map((m) => ({
+            [this.memberName(m)]: {
+              sql: this.sqlForMember(m),
+              type: m.type ?? m.types[0],
+              title: this.memberTitle(m),
+              [this.options.snakeCase ? 'primary_key' : 'primaryKey']:
+                  m.isPrimaryKey ? true : undefined,
+            },
+          }))
+          .reduce((a, b) => ({ ...a, ...b }), {})
+      },
       measures: tableSchema.measures
         .map((m) => ({
           [this.memberName(m)]: {
@@ -172,21 +211,24 @@ export abstract class BaseSchemaFormatter {
 
       ...(this.options.snakeCase
         ? Object.fromEntries(
-          Object.entries(contextProps).map(([key, value]) => [toSnakeCase(key), value])
+          Object.entries(contextProps).map(([key, value]) => [
+            toSnakeCase(key),
+            value,
+          ])
         )
         : contextProps),
 
-      [this.options.snakeCase ? 'pre_aggregations' : 'preAggregations']: new ValueWithComments(
-        null,
-        [
+      [this.options.snakeCase ? 'pre_aggregations' : 'preAggregations']:
+        new ValueWithComments(null, [
           'Pre-aggregation definitions go here.',
           'Learn more in the documentation: https://cube.dev/docs/caching/pre-aggregations/getting-started',
-        ]
-      ),
+        ]),
     };
   }
 
-  protected schemaForTablesByCubeDescriptors(cubeDescriptors: CubeDescriptor[]) {
+  protected schemaForTablesByCubeDescriptors(
+    cubeDescriptors: CubeDescriptor[]
+  ) {
     const tableNames = cubeDescriptors.map(({ tableName }) => tableName);
     const generatedSchemaForTables = this.scaffoldingSchema.generateForTables(
       tableNames.map((n) => this.scaffoldingSchema.resolveTableName(n))
@@ -212,10 +254,16 @@ export abstract class BaseSchemaFormatter {
         }
       );
 
+      let compositePrimaryKey: CubeDescriptorMember[] | undefined = [];
+      if (descriptor.shouldGeneratePrimaryKey) {
+        compositePrimaryKey = cubeMembers.dimensions.filter(d => !d.isPrimaryKey && d.type !== 'time');
+      }
+
       return {
         ...generatedDescriptor,
         ...descriptor,
         ...cubeMembers,
+        compositePrimaryKey
       };
     });
   }

diff --git a/packages/cubejs-schema-compiler/test/unit/scaffolding-template.test.ts b/packages/cubejs-schema-compiler/test/unit/scaffolding-template.test.ts
@@ -1,14 +1,17 @@
+import { MemberType } from '../../src';
 import {
   ScaffoldingTemplate,
   SchemaFormat,
 } from '../../src/scaffolding/ScaffoldingTemplate';
 
 const driver = {
   quoteIdentifier: (name) => `"${name}"`,
+  concatStringsSql: (strings: string[]) => strings.join(' || '),
 };
 
 const mySqlDriver = {
   quoteIdentifier: (name) => `\`${name}\``,
+  concatStringsSql: (strings: string[]) => `CONCAT(${strings.join(', ')})`
 };
 
 const bigQueryDriver = {
@@ -587,6 +590,119 @@ describe('ScaffoldingTemplate', () => {
   });
 
   describe('Yaml formatter', () => {
+    it('generates composite primary key for YAML model', () => {
+      const template = new ScaffoldingTemplate(dbSchema, driver, {
+        format: SchemaFormat.Yaml,
+        snakeCase: true
+      });
+
+      const files = template.generateFilesByCubeDescriptors([
+        {
+          cube: 'orders',
+          tableName: 'public.orders',
+          table: 'orders',
+          schema: 'public',
+          joins: [],
+          shouldGeneratePrimaryKey: true,
+          members: [
+            {
+              memberType: MemberType.Dimension,
+              name: 'number',
+              title: 'Number',
+              types: [
+                'sum',
+                'avg',
+                'min',
+                'max'
+              ],
+              included: true,
+              type: 'number'
+            },
+            {
+              memberType: MemberType.Dimension,
+              name: 'status',
+              title: 'Status',
+              types: [
+                'string'
+              ],
+              isPrimaryKey: false,
+              type: 'string',
+              included: true
+            },
+            {
+              memberType: MemberType.Dimension,
+              name: 'created_at',
+              title: 'Created At',
+              types: [
+                'number'
+              ],
+              type: 'time',
+              included: true
+            },
+          ],
+        },
+      ]);
+
+      expect(files[0].content).toContain('sql: "{CUBE}.number || \'-\' || {CUBE}.status"');
+    });
+
+    it('generates composite primary key for JS model', () => {
+      const template = new ScaffoldingTemplate(dbSchema, driver, {
+        format: SchemaFormat.JavaScript,
+        snakeCase: true
+      });
+
+      const files = template.generateFilesByCubeDescriptors([
+        {
+          cube: 'orders',
+          tableName: 'public.orders',
+          table: 'orders',
+          schema: 'public',
+          joins: [],
+          shouldGeneratePrimaryKey: true,
+          members: [
+            {
+              memberType: MemberType.Dimension,
+              name: 'number',
+              title: 'Number',
+              types: [
+                'sum',
+                'avg',
+                'min',
+                'max'
+              ],
+              included: true,
+              type: 'number'
+            },
+            {
+              memberType: MemberType.Dimension,
+              name: 'status',
+              title: 'Status',
+              types: [
+                'string'
+              ],
+              isPrimaryKey: false,
+              type: 'string',
+              included: true
+            },
+            {
+              memberType: MemberType.Dimension,
+              name: 'created_at',
+              title: 'Created At',
+              types: [
+                'number'
+              ],
+              type: 'time',
+              included: true
+            },
+          ],
+        },
+      ]);
+
+      // eslint-disable-next-line
+      expect(files[0].content).toContain('sql: `${CUBE}.number || \'-\' || ${CUBE}.status`');
+    });
+
     it('generates schema for base driver', () => {
       const template = new ScaffoldingTemplate(dbSchema, driver, {
         format: SchemaFormat.Yaml,