feat(schema-compiler): Allow passing function that returns an array o…

…f time dimensions as pre-aggregation timeDimensions property
cube-js · Jan 23, 2025 · 293f833 · 293f833
1 parent 4583816
commit 293f833
Show file tree

Hide file tree

Showing 4 changed files with 183 additions and 5 deletions.
diff --git a/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts b/packages/cubejs-schema-compiler/src/compiler/CubeEvaluator.ts
@@ -696,14 +696,22 @@ export class CubeEvaluator extends CubeSymbols {
         dimension: this.evaluateReferences(cube, aggregation.timeDimensionReference),
         granularity: aggregation.granularity
       });
-    } else if (aggregation.timeDimensionReferences) {
+    } else if (Array.isArray(aggregation.timeDimensionReferences)) {
       // eslint-disable-next-line guard-for-in
       for (const timeDimensionReference of aggregation.timeDimensionReferences) {
         timeDimensions.push({
           dimension: this.evaluateReferences(cube, timeDimensionReference.dimension),
           granularity: timeDimensionReference.granularity
         });
       }
+    } else if (aggregation.timeDimensionReferences) {
+      const evaluatedRefs: any[] = this.evaluateReferences(cube, aggregation.timeDimensionReferences, { returnRaw: true });
+      for (const timeDimensionReference of evaluatedRefs) {
+        timeDimensions.push({
+          dimension: timeDimensionReference.dimension.toString(),
+          granularity: timeDimensionReference.granularity
+        });
+      }
     }
 
     return {

diff --git a/packages/cubejs-schema-compiler/src/compiler/CubeSymbols.js b/packages/cubejs-schema-compiler/src/compiler/CubeSymbols.js
@@ -515,6 +515,11 @@ export class CubeSymbols {
       cubeAliasFn: (cube) => cubeEvaluator.pathFromArray(fullPath(cubeEvaluator.joinHints(), [cube])),
       collectJoinHints: options.collectJoinHints,
     });
+
+    if (options.returnRaw) {
+      return arrayOrSingle;
+    }
+
     if (!Array.isArray(arrayOrSingle)) {
       return arrayOrSingle.toString();
     }

diff --git a/packages/cubejs-schema-compiler/src/compiler/CubeValidator.ts b/packages/cubejs-schema-compiler/src/compiler/CubeValidator.ts
@@ -473,10 +473,13 @@ const RollUpSchema = condition(
       // Rollup with multiple time dimensions
       inherit(BasePreAggregation, {
         type: Joi.any().valid('rollup').required(),
-        timeDimensions: Joi.array().items(Joi.object().keys({
-          dimension: Joi.func(),
-          granularity: GranularitySchema,
-        })),
+        timeDimensions: Joi.alternatives().try(
+          Joi.array().items(Joi.object().keys({
+            dimension: Joi.func(),
+            granularity: GranularitySchema,
+          })),
+          Joi.func(),
+        ),
         allowNonStrictDateRangeMatch: Joi.bool(),
         measures: Joi.func(),
         dimensions: Joi.func(),

diff --git a/packages/cubejs-schema-compiler/test/unit/pre-aggregations.test.ts b/packages/cubejs-schema-compiler/test/unit/pre-aggregations.test.ts
@@ -92,6 +92,168 @@ describe('pre-aggregations', () => {
     expect(cubeEvaluator.cubeFromPath('Orders').preAggregations.ordersRollupJoin.scheduledRefresh).toEqual(undefined);
   });
 
+  it('Rollup with pre-agg with hardcoded multiple time dimensions', async () => {
+    const { compiler, cubeEvaluator, joinGraph } = prepareCompiler(
+      `
+        cube(\`Users\`, {
+          sql: \`SELECT * FROM public.users\`,
+
+          preAggregations: {
+            staticMultiple: {
+              dimensions: [CUBE.status],
+              measures: [CUBE.count],
+              timeDimensions: [
+                { dimension: CUBE.createdAt, granularity: \`day\` },
+                { dimension: CUBE.modifiedAt, granularity: \`day\` },
+              ]
+            }
+          },
+
+          measures: {
+            count: {
+              type: \`count\`,
+            },
+          },
+
+          dimensions: {
+            id: {
+              sql: \`id\`,
+              type: \`string\`,
+              primaryKey: true,
+            },
+
+            name: {
+              sql: \`name\`,
+              type: \`string\`,
+            },
+
+            userId: {
+              sql: \`user_id\`,
+              type: \`number\`,
+            },
+            status: {
+              sql: \`status\`,
+              type: \`string\`,
+            },
+
+            createdAt: {
+              type: \`time\`,
+              sql: \`created_at\`
+            },
+
+            modifiedAt: {
+              type: \`time\`,
+              sql: \`modified_at\`
+            }
+          },
+        });
+      `
+    );
+
+    await compiler.compile();
+
+    const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
+      dimensions: ['Users.status'],
+      measures: ['Users.count'],
+      timeDimensions: [{
+        dimension: 'Users.createdAt',
+        dateRange: ['2023-01-20', '2024-01-20'],
+        granularity: 'day'
+      }]
+    });
+
+    const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription();
+
+    const queryAndParams = query.buildSqlAndParams();
+    console.log(queryAndParams);
+    expect(queryAndParams[0].includes('undefined')).toBeFalsy();
+    expect(queryAndParams[0].includes('pre_aggregations')).toBeTruthy();
+
+    expect(preAggregationsDescription.length).toEqual(1);
+    expect(preAggregationsDescription[0].preAggregationId).toEqual('Users.staticMultiple');
+  });
+
+  it('Rollup with pre-agg with dynamic multiple time dimensions', async () => {
+    const { compiler, cubeEvaluator, joinGraph } = prepareCompiler(
+      `
+        cube(\`Users\`, {
+          sql: \`SELECT * FROM public.users\`,
+
+          preAggregations: {
+            dynamicMultiple: {
+              dimensions: [CUBE.status],
+              measures: [CUBE.count],
+              timeDimensions: (CUBE) => [
+                {dimension: CUBE.createdAt, granularity: 'day'},
+                {dimension: CUBE.modifiedAt, granularity: 'day'},
+              ]
+            },
+          },
+
+          measures: {
+            count: {
+              type: \`count\`,
+            },
+          },
+
+          dimensions: {
+            id: {
+              sql: \`id\`,
+              type: \`string\`,
+              primaryKey: true,
+            },
+
+            name: {
+              sql: \`name\`,
+              type: \`string\`,
+            },
+
+            userId: {
+              sql: \`user_id\`,
+              type: \`number\`,
+            },
+            status: {
+              sql: \`status\`,
+              type: \`string\`,
+            },
+
+            createdAt: {
+              type: \`time\`,
+              sql: \`created_at\`
+            },
+
+            modifiedAt: {
+              type: \`time\`,
+              sql: \`modified_at\`
+            }
+          },
+        });
+      `
+    );
+
+    await compiler.compile();
+
+    const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
+      dimensions: ['Users.status'],
+      measures: ['Users.count'],
+      timeDimensions: [{
+        dimension: 'Users.createdAt',
+        dateRange: ['2023-01-20', '2024-01-20'],
+        granularity: 'day'
+      }]
+    });
+
+    const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription();
+
+    const queryAndParams = query.buildSqlAndParams();
+    console.log(queryAndParams);
+    expect(queryAndParams[0].includes('undefined')).toBeFalsy();
+    expect(queryAndParams[0].includes('pre_aggregations')).toBeTruthy();
+
+    expect(preAggregationsDescription.length).toEqual(1);
+    expect(preAggregationsDescription[0].preAggregationId).toEqual('Users.dynamicMultiple');
+  });
+
   it('query rollupLambda', async () => {
     const { compiler, cubeEvaluator, joinGraph } = prepareCompiler(
       `