feat(cubesql): DATE_TRUNC support for SQL push down (cube-js#7132)

ConsultingMD · Sep 14, 2023 · ae80eb1 · ae80eb1
1 parent 30356d9
commit ae80eb1
Show file tree

Hide file tree

Showing 12 changed files with 64 additions and 4 deletions.
diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts
@@ -1650,7 +1650,8 @@ const variables: Record<string, (...args: any) => any> = {
     .asInt(),
   maxSourceRowLimit: () => get('CUBEJS_MAX_SOURCE_ROW_LIMIT')
     .default(200000)
-    .asInt()
+    .asInt(),
+  convertTzForRawTimeDimension: () => get('CUBESQL_SQL_PUSH_DOWN').default('false').asBoolStrict(),
 };
 
 type Vars = typeof variables;

diff --git a/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts b/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts
@@ -114,4 +114,10 @@ export class DatabricksQuery extends BaseQuery {
   public defaultRefreshKeyRenewalThreshold() {
     return 120;
   }
+
+  public sqlTemplates() {
+    const templates = super.sqlTemplates();
+    templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
+    return templates;
+  }
 }
diff --git a/packages/cubejs-duckdb-driver/src/DuckDBQuery.ts b/packages/cubejs-duckdb-driver/src/DuckDBQuery.ts
@@ -30,4 +30,10 @@ export class DuckDBQuery extends BaseQuery {
   public countDistinctApprox(sql: string) {
     return `approx_count_distinct(${sql})`;
   }
+
+  public sqlTemplates() {
+    const templates = super.sqlTemplates();
+    templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
+    return templates;
+  }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseDimension.js b/packages/cubejs-schema-compiler/src/adapter/BaseDimension.js
@@ -24,12 +24,22 @@ export class BaseDimension {
 
   dimensionSql() {
     if (this.expression) {
-      return this.query.evaluateSymbolSql(this.expressionCubeName, this.expressionName, this.definition(), 'dimension');
+      return this.convertTzForRawTimeDimensionIfNeeded(() => this.query.evaluateSymbolSql(this.expressionCubeName, this.expressionName, this.definition(), 'dimension'));
     }
     if (this.query.cubeEvaluator.isSegment(this.dimension)) {
       return this.query.wrapSegmentForDimensionSelect(this.query.dimensionSql(this));
     }
-    return this.query.dimensionSql(this);
+    return this.convertTzForRawTimeDimensionIfNeeded(() => this.query.dimensionSql(this));
+  }
+
+  convertTzForRawTimeDimensionIfNeeded(sql) {
+    if (this.query.options.convertTzForRawTimeDimension) {
+      return this.query.evaluateSymbolSqlWithContext(sql, {
+        convertTzForRawTimeDimension: true
+      });
+    } else {
+      return sql();
+    }
   }
 
   sqlDefinition() {

diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseFilter.js b/packages/cubejs-schema-compiler/src/adapter/BaseFilter.js
@@ -32,6 +32,10 @@ export class BaseFilter extends BaseDimension {
     return this.conditionSql(this.measure ? this.query.measureSql(this) : this.query.dimensionSql(this));
   }
 
+  convertTzForRawTimeDimensionIfNeeded(sql) {
+    return sql();
+  }
+
   // Evaluates filters on measures to whole where statement in query
   // It used in drill downs
   measureFilterToWhere() {

diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js
@@ -210,6 +210,7 @@ class BaseQuery {
       ungrouped: this.options.ungrouped,
       memberToAlias: this.options.memberToAlias,
       expressionParams: this.options.expressionParams,
+      convertTzForRawTimeDimension: this.options.convertTzForRawTimeDimension,
     });
     this.timezone = this.options.timezone;
     this.rowLimit = this.options.rowLimit;
@@ -1796,7 +1797,15 @@ class BaseQuery {
           this.autoPrefixAndEvaluateSql(cubeName, symbol.longitude.sql)
         ]);
       } else {
-        return this.autoPrefixAndEvaluateSql(cubeName, symbol.sql);
+        let res = this.autoPrefixAndEvaluateSql(cubeName, symbol.sql);
+        if (this.safeEvaluateSymbolContext().convertTzForRawTimeDimension &&
+          !memberExpressionType &&
+          symbol.type === 'time' &&
+          this.cubeEvaluator.byPathAnyType(memberPathArray).ownedByCube
+        ) {
+          res = this.convertTz(res);
+        }
+        return res;
       }
     } else if (type === 'segment') {
       if ((this.safeEvaluateSymbolContext().renderedReference || {})[memberPath]) {
@@ -2399,6 +2408,10 @@ class BaseQuery {
     return false;
   }
 
+  /**
+   * @public
+   * @returns {any}
+   */
   sqlTemplates() {
     return {
       functions: {

diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseTimeDimension.js b/packages/cubejs-schema-compiler/src/adapter/BaseTimeDimension.js
@@ -77,6 +77,10 @@ export class BaseTimeDimension extends BaseFilter {
     return this.query.timeGroupedColumn(granularity, this.convertedToTz());
   }
 
+  convertTzForRawTimeDimensionIfNeeded(sql) {
+    return sql();
+  }
+
   convertedToTz() {
     return this.query.convertTz(this.query.dimensionSql(this));
   }

diff --git a/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.js b/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.js
@@ -48,6 +48,7 @@ export class PostgresQuery extends BaseQuery {
     const templates = super.sqlTemplates();
     // eslint-disable-next-line no-template-curly-in-string
     templates.params.param = '${{ param_index + 1 }}';
+    templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
     return templates;
   }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.js b/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.js
@@ -108,4 +108,10 @@ export class PrestodbQuery extends BaseQuery {
 
     return `${offsetClause}${limitClause}`;
   }
+
+  sqlTemplates() {
+    const templates = super.sqlTemplates();
+    templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
+    return templates;
+  }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.js b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.js
@@ -49,4 +49,10 @@ export class SnowflakeQuery extends BaseQuery {
   countDistinctApprox(sql) {
     return `APPROX_COUNT_DISTINCT(${sql})`;
   }
+
+  sqlTemplates() {
+    const templates = super.sqlTemplates();
+    templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})';
+    return templates;
+  }
 }
diff --git a/packages/cubejs-server-core/src/core/CompilerApi.js b/packages/cubejs-server-core/src/core/CompilerApi.js
@@ -20,6 +20,7 @@ export class CompilerApi {
     this.logger = this.options.logger;
     this.preAggregationsSchema = this.options.preAggregationsSchema;
     this.allowUngroupedWithoutPrimaryKey = this.options.allowUngroupedWithoutPrimaryKey;
+    this.convertTzForRawTimeDimension = this.options.convertTzForRawTimeDimension;
     this.schemaVersion = this.options.schemaVersion;
     this.compileContext = options.compileContext;
     this.allowJsDuplicatePropsInSchema = options.allowJsDuplicatePropsInSchema;
@@ -206,6 +207,7 @@ export class CompilerApi {
         externalDbType: this.options.externalDbType,
         preAggregationsSchema: this.preAggregationsSchema,
         allowUngroupedWithoutPrimaryKey: this.allowUngroupedWithoutPrimaryKey,
+        convertTzForRawTimeDimension: this.convertTzForRawTimeDimension,
         queryFactory: this.queryFactory,
       }
     );

diff --git a/packages/cubejs-server-core/src/core/server.ts b/packages/cubejs-server-core/src/core/server.ts
@@ -675,6 +675,7 @@ export class CubejsServerCore {
         allowUngroupedWithoutPrimaryKey:
             this.options.allowUngroupedWithoutPrimaryKey ||
             getEnv('allowUngroupedWithoutPrimaryKey'),
+        convertTzForRawTimeDimension: getEnv('convertTzForRawTimeDimension'),
         compileContext: options.context,
         dialectClass: options.dialectClass,
         externalDialectClass: options.externalDialectClass,