fix: Cross data source rollupLambda uses incorrect refreshKey (#6288)

* fix: Cross data source `rollupLambda` uses incorrect refreshKey

* fix(ksql-driver): Missing quote identifier

* fix(ksql-driver): Missing quote identifier

Author: paveltiunov

packages/cubejs-ksql-driver/src/KsqlQuery.ts

@@ -69,6 +69,9 @@ export class KsqlQuery extends BaseQuery {
   }
 
   public groupByClause() {
+    if (this.ungrouped) {
+      return '';
+    }
     const dimensionsForSelect: any[] = this.dimensionsForSelect();
     const dimensionColumns = dimensionsForSelect.map(s => s.selectColumns() && s.dimensionSql())
       .reduce((a, b) => a.concat(b), [])
@@ -99,7 +102,12 @@ export class KsqlQuery extends BaseQuery {
 
   public preAggregationReadOnly(cube: string, preAggregation: any) {
     const [sql] = this.preAggregationSql(cube, preAggregation);
-    return preAggregation.type === 'originalSql' && Boolean(KsqlQuery.extractTableFromSimpleSelectAsteriskQuery(sql));
+    return preAggregation.type === 'originalSql' && Boolean(KsqlQuery.extractTableFromSimpleSelectAsteriskQuery(sql)) ||
+      preAggregation.type === 'rollup' && !!this.dimensionsForSelect().find(d => d.definition().primaryKey);
+  }
+
+  public preAggregationAllowUngroupingWithPrimaryKey(_cube: any, _preAggregation: any) {
+    return true;
   }
 
   public static extractTableFromSimpleSelectAsteriskQuery(sql: string) {

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -338,9 +338,6 @@ class BaseQuery {
           throw new UserError(`Ungrouped query requires primary keys to be present in dimensions: ${missingPrimaryKeys.map(k => `'${k}'`).join(', ')}. Pass allowUngroupedWithoutPrimaryKey option to disable this check.`);
         }
       }
-      if (this.measures.length) {
-        throw new UserError('Measures aren\'t allowed in ungrouped query');
-      }
       if (this.measureFilters.length) {
         throw new UserError('Measure filters aren\'t allowed in ungrouped query');
       }
@@ -1956,6 +1953,13 @@ class BaseQuery {
     ) {
       return evaluateSql === '*' ? '1' : evaluateSql;
     }
+    if (this.ungrouped) {
+      if (symbol.type === 'count' || symbol.type === 'countDistinct' || symbol.type === 'countDistinctApprox') {
+        return '1';
+      } else {
+        return evaluateSql;
+      }
+    }
     if ((this.safeEvaluateSymbolContext().ungroupedAliases || {})[measurePath]) {
       evaluateSql = (this.safeEvaluateSymbolContext().ungroupedAliases || {})[measurePath];
     }
@@ -2368,6 +2372,10 @@ class BaseQuery {
     return false;
   }
 
+  preAggregationAllowUngroupingWithPrimaryKey(_cube, _preAggregation) {
+    return false;
+  }
+
   // eslint-disable-next-line consistent-return
   preAggregationQueryForSqlEvaluation(cube, preAggregation) {
     if (preAggregation.type === 'autoRollup') {

packages/cubejs-schema-compiler/src/adapter/PreAggregations.js

@@ -148,7 +148,8 @@ export class PreAggregations {
 
     const tableName = this.preAggregationTableName(cube, preAggregationName, preAggregation);
     const invalidateKeyQueries = this.query.preAggregationInvalidateKeyQueries(cube, preAggregation);
-    const partitionInvalidateKeyQueries = this.query.partitionInvalidateKeyQueries && this.query.partitionInvalidateKeyQueries(cube, preAggregation);
+    const queryForSqlEvaluation = this.query.preAggregationQueryForSqlEvaluation(cube, preAggregation);
+    const partitionInvalidateKeyQueries = queryForSqlEvaluation.partitionInvalidateKeyQueries && queryForSqlEvaluation.partitionInvalidateKeyQueries(cube, preAggregation);
 
     const matchedTimeDimension =
       preAggregation.partitionGranularity &&
@@ -161,11 +162,10 @@ export class PreAggregations {
         td.isDateOperator() &&
         td.camelizeOperator === 'inDateRange' // TODO support all date operators
     );
-    const queryForSqlEvaluation = this.query.preAggregationQueryForSqlEvaluation(cube, preAggregation);
 
     const uniqueKeyColumnsDefault = () => null;
     const uniqueKeyColumns = ({
-      rollup: () => this.query.preAggregationQueryForSqlEvaluation(cube, preAggregation).dimensionColumns(),
+      rollup: () => queryForSqlEvaluation.dimensionColumns(),
       originalSql: () => preAggregation.uniqueKeyColumns || null
     }[preAggregation.type] || uniqueKeyColumnsDefault)();
 
@@ -180,18 +180,18 @@ export class PreAggregations {
       partitionInvalidateKeyQueries,
       type: preAggregation.type,
       external: preAggregation.external,
-      previewSql: this.query.preAggregationPreviewSql(tableName),
-      preAggregationsSchema: this.query.preAggregationSchema(),
-      loadSql: this.query.preAggregationLoadSql(cube, preAggregation, tableName),
-      sql: this.query.preAggregationSql(cube, preAggregation),
+      previewSql: queryForSqlEvaluation.preAggregationPreviewSql(tableName),
+      preAggregationsSchema: queryForSqlEvaluation.preAggregationSchema(),
+      loadSql: queryForSqlEvaluation.preAggregationLoadSql(cube, preAggregation, tableName),
+      sql: queryForSqlEvaluation.preAggregationSql(cube, preAggregation),
       uniqueKeyColumns,
       aggregationsColumns,
       dataSource: queryForSqlEvaluation.dataSource,
       // in fact we can reference preAggregation.granularity however accessing timeDimensions is more strict and consistent
       granularity: references.timeDimensions[0]?.granularity,
       partitionGranularity: preAggregation.partitionGranularity,
       updateWindowSeconds: preAggregation.refreshKey && preAggregation.refreshKey.updateWindow &&
-        this.query.parseSecondDuration(preAggregation.refreshKey.updateWindow),
+        queryForSqlEvaluation.parseSecondDuration(preAggregation.refreshKey.updateWindow),
       preAggregationStartEndQueries:
         (preAggregation.partitionGranularity || references.timeDimensions[0]?.granularity) &&
         this.refreshRangeQuery().preAggregationStartEndQueries(cube, preAggregation),
@@ -208,7 +208,7 @@ export class PreAggregations {
             const indexName = this.preAggregationTableName(cube, `${foundPreAggregation.sqlAlias || preAggregationName}_${index}`, preAggregation, true);
             return {
               indexName,
-              sql: this.query.indexSql(
+              sql: queryForSqlEvaluation.indexSql(
                 cube,
                 preAggregation,
                 preAggregation.indexes[index],
@@ -226,11 +226,11 @@ export class PreAggregations {
             return {
               indexName,
               type: preAggregation.indexes[index].type,
-              columns: this.query.evaluateIndexColumns(cube, preAggregation.indexes[index])
+              columns: queryForSqlEvaluation.evaluateIndexColumns(cube, preAggregation.indexes[index])
             };
           }
         ),
-      readOnly: preAggregation.readOnly || this.query.preAggregationReadOnly(cube, preAggregation),
+      readOnly: preAggregation.readOnly || queryForSqlEvaluation.preAggregationReadOnly(cube, preAggregation),
       streamOffset: preAggregation.streamOffset,
       unionWithSourceData: preAggregation.unionWithSourceData,
       rollupLambdaId: preAggregation.rollupLambdaId,
@@ -978,6 +978,7 @@ export class PreAggregations {
 
   rollupPreAggregationQuery(cube, aggregation) {
     const references = this.evaluateAllReferences(cube, aggregation);
+    const cubeQuery = this.query.newSubQueryForCube(cube, {});
     return this.query.newSubQueryForCube(
       cube,
       {
@@ -987,6 +988,8 @@ export class PreAggregations {
         timeDimensions: this.mergePartitionTimeDimensions(references, aggregation.partitionTimeDimensions),
         preAggregationQuery: true,
         useOriginalSqlPreAggregationsInPreAggregation: aggregation.useOriginalSqlPreAggregations,
+        ungrouped: cubeQuery.preAggregationAllowUngroupingWithPrimaryKey(cube, aggregation) &&
+          !!references.dimensions.find(d => this.query.cubeEvaluator.dimensionByPath(d).primaryKey)
       }
     );
   }

packages/cubejs-schema-compiler/test/integration/postgres/pre-aggregations.test.ts

@@ -410,6 +410,46 @@ describe('PreAggregations', () => {
         }
       },
     });
+    
+    cube('LambdaVisitors', {
+      extends: visitors,
+      
+      preAggregations: {
+        partitionedLambda: {
+          type: 'rollupLambda',
+          rollups: [partitioned, RealTimeLambdaVisitors.partitioned]
+        },
+        partitioned: {
+          type: 'rollup',
+          measures: [count],
+          dimensions: [id, source],
+          timeDimension: createdAt,
+          granularity: 'day',
+          partitionGranularity: 'month',
+          refreshKey: {
+            every: '1 hour',
+            incremental: true,
+            updateWindow: '1 day'
+          }
+        }
+      }
+    });
+    
+    cube('RealTimeLambdaVisitors', {
+      dataSource: 'ksql',
+      extends: visitors,
+      
+      preAggregations: {
+        partitioned: {
+          type: 'rollup',
+          measures: [count],
+          dimensions: [id, source],
+          timeDimension: createdAt,
+          granularity: 'day',
+          partitionGranularity: 'day'
+        }
+      }
+    });
     `);
 
   it('simple pre-aggregation', () => compiler.compile().then(() => {
@@ -1629,4 +1669,46 @@ describe('PreAggregations', () => {
       );
     });
   }));
+
+  it('lambda cross data source refresh key and ungrouped', () => compiler.compile().then(() => {
+    const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
+      measures: [
+        'LambdaVisitors.count'
+      ],
+      dimensions: [
+        'LambdaVisitors.source'
+      ],
+      timeDimensions: [{
+        dimension: 'LambdaVisitors.createdAt',
+        granularity: 'day',
+        dateRange: ['2017-01-01', '2017-01-25']
+      }],
+      timezone: 'America/Los_Angeles',
+      order: [{
+        id: 'LambdaVisitors.createdAt'
+      }],
+      preAggregationsSchema: '',
+      queryFactory: {
+        createQuery: (cube, compilers, options) => {
+          if (cube === 'RealTimeLambdaVisitors') {
+            // eslint-disable-next-line global-require
+            const { KsqlQuery } = require('../../../../../cubejs-ksql-driver');
+            return new KsqlQuery(compilers, options);
+          } else {
+            return new PostgresQuery(compilers, options);
+          }
+        }
+      }
+    });
+
+    const queryAndParams = query.buildSqlAndParams();
+    console.log(queryAndParams);
+    const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription();
+    console.log(JSON.stringify(preAggregationsDescription, null, 2));
+    const { partitionInvalidateKeyQueries, loadSql } = preAggregationsDescription.find(p => p.preAggregationId === 'RealTimeLambdaVisitors.partitioned');
+
+    expect(partitionInvalidateKeyQueries).toStrictEqual([]);
+    expect(loadSql[0]).not.toMatch(/GROUP BY/);
+    expect(loadSql[0]).toMatch(/1 `real_time_lambda_visitors__count`/);
+  }));
 });