From bb27698687ec09fad2e9d6ea96b63c75ba4ad249 Mon Sep 17 00:00:00 2001
From: Pavel Tiunov <pavel.tiunov@gmail.com>
Date: Mon, 12 Feb 2024 23:32:40 -0800
Subject: [PATCH] fix(cubesql): SQL push down can be incorrectly routed to Cube
 Store pre-aggregations

---
 packages/cubejs-api-gateway/src/gateway.ts     | 18 +++++++++++++++---
 packages/cubejs-api-gateway/src/sql-server.ts  |  1 +
 .../cubejs-api-gateway/src/types/request.ts    |  1 +
 .../src/adapter/BaseQuery.js                   |  8 ++++++--
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts
index 0899d470f1a1e..09ccab15a7d3b 100644
--- a/packages/cubejs-api-gateway/src/gateway.ts
+++ b/packages/cubejs-api-gateway/src/gateway.ts
@@ -1172,7 +1172,16 @@ class ApiGateway {
     return [queryType, normalizedQueries];
   }
 
-  public async sql({ query, context, res, memberToAlias, exportAnnotatedSql, memberExpressions, expressionParams }: QueryRequest) {
+  public async sql({
+    query,
+    context,
+    res,
+    memberToAlias,
+    exportAnnotatedSql,
+    memberExpressions,
+    expressionParams,
+    disableExternalPreAggregations
+  }: QueryRequest) {
     const requestStarted = new Date();
 
     try {
@@ -1188,7 +1197,7 @@ class ApiGateway {
 
       const sqlQueries = await Promise.all<any>(
         normalizedQueries.map(async (normalizedQuery) => (await this.getCompilerApi(context)).getSql(
-          this.coerceForSqlQuery({ ...normalizedQuery, memberToAlias, expressionParams }, context),
+          this.coerceForSqlQuery({ ...normalizedQuery, memberToAlias, expressionParams, disableExternalPreAggregations }, context),
           {
             includeDebugInfo: getEnv('devMode') || context.signedWithPlaygroundAuthSecret,
             exportAnnotatedSql,
@@ -1708,7 +1717,10 @@ class ApiGateway {
       metaConfigResult = this.filterVisibleItemsInMeta(context, metaConfigResult);
 
       const sqlQueries = await this
-        .getSqlQueriesInternal(context, normalizedQueries);
+        .getSqlQueriesInternal(
+          context,
+          normalizedQueries.map(q => ({ ...q, disableExternalPreAggregations: request.sqlQuery }))
+        );
 
       let results;
 
diff --git a/packages/cubejs-api-gateway/src/sql-server.ts b/packages/cubejs-api-gateway/src/sql-server.ts
index b5f914e185761..efdb00f50367e 100644
--- a/packages/cubejs-api-gateway/src/sql-server.ts
+++ b/packages/cubejs-api-gateway/src/sql-server.ts
@@ -157,6 +157,7 @@ export class SQLServer {
               expressionParams,
               exportAnnotatedSql: true,
               memberExpressions: true,
+              disableExternalPreAggregations: true,
               queryType: 'multi',
               context,
               res: (message) => {
diff --git a/packages/cubejs-api-gateway/src/types/request.ts b/packages/cubejs-api-gateway/src/types/request.ts
index 1b9394f193db7..c6c2a881333f9 100644
--- a/packages/cubejs-api-gateway/src/types/request.ts
+++ b/packages/cubejs-api-gateway/src/types/request.ts
@@ -126,6 +126,7 @@ type QueryRequest = BaseRequest & {
   expressionParams?: string[];
   exportAnnotatedSql?: boolean;
   memberExpressions?: boolean;
+  disableExternalPreAggregations?: boolean;
 };
 
 type SqlApiRequest = BaseRequest & {
diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js
index c224c5dac4c99..5105d498b9aef 100644
--- a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js
+++ b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js
@@ -202,6 +202,7 @@ export class BaseQuery {
       className: this.constructor.name,
       externalClassName: this.options.externalQueryClass && this.options.externalQueryClass.name,
       preAggregationQuery: this.options.preAggregationQuery,
+      disableExternalPreAggregations: this.options.disableExternalPreAggregations,
       useOriginalSqlPreAggregationsInPreAggregation: this.options.useOriginalSqlPreAggregationsInPreAggregation,
       cubeLatticeCache: this.options.cubeLatticeCache, // TODO too heavy for key
       historyQueries: this.options.historyQueries, // TODO too heavy for key
@@ -467,6 +468,9 @@ export class BaseQuery {
     if (!this.options.preAggregationQuery && !this.ungrouped) {
       preAggForQuery =
         this.preAggregations.findPreAggregationForQuery();
+      if (this.options.disableExternalPreAggregations && preAggForQuery.preAggregation.external) {
+        preAggForQuery = undefined;
+      }
     }
     if (preAggForQuery) {
       const {
@@ -537,7 +541,7 @@ export class BaseQuery {
   }
 
   externalPreAggregationQuery() {
-    if (!this.options.preAggregationQuery && this.externalQueryClass) {
+    if (!this.options.preAggregationQuery && !this.options.disableExternalPreAggregations && this.externalQueryClass) {
       const preAggregationForQuery = this.preAggregations.findPreAggregationForQuery();
       if (preAggregationForQuery && preAggregationForQuery.preAggregation.external) {
         return true;
@@ -556,7 +560,7 @@ export class BaseQuery {
    * @returns {Array<string>}
    */
   buildSqlAndParams(exportAnnotatedSql) {
-    if (!this.options.preAggregationQuery && this.externalQueryClass) {
+    if (!this.options.preAggregationQuery && !this.options.disableExternalPreAggregations && this.externalQueryClass) {
       if (this.externalPreAggregationQuery()) { // TODO performance
         return this.externalQuery().buildSqlAndParams(exportAnnotatedSql);
       }