feat(cubesql): Whole SQL query push down to data sources (#6629)

* feat(cubesql): Whole SQL query push down to data sources

* Fix clippy

* Fix clippy

* Fix cost function

* Add sqlGenerators mock

* Fix tests

* SQL Generation structure

* Fix tests

* SQL transport API

* Query sent to load method but member re-projection is required

* Simple querying works

* Move to minijinja

* Clippy and formatter

* Linter

* Update chrono

* Restructure templates

* Simple projection rules

* Introduce SQL push down flag and make sure rules still work with it disabled

* Fix projection positive feedback loop and implement coalesce function as an example of function pushdown

* Parameters implementation

* Missing jsdoc param

* SQL API load support

* SQL API load streaming support: fix unit tests

* SQL API load streaming support: fix driver tests

* SQL API load streaming support: update mssql snapshot

* Revert "SQL API load streaming support: update mssql snapshot"

This reverts commit 91ed17e.

* SQL API load streaming support: update schema

* Bump rust toolchain to address tarpaulin issue

Author: paveltiunov

.github/workflows/rust-cubesql.yml

@@ -61,7 +61,7 @@ jobs:
       - name: Install Rust
         uses: actions-rs/toolchain@v1
         with:
-          toolchain: nightly-2022-03-22
+          toolchain: nightly-2022-06-22
           override: true
           components: rustfmt
       - uses: Swatinem/rust-cache@v2
@@ -80,6 +80,7 @@ jobs:
           CUBESQL_TESTING_CUBE_TOKEN: ${{ secrets.CUBESQL_TESTING_CUBE_TOKEN }}
           CUBESQL_TESTING_CUBE_URL: ${{ secrets.CUBESQL_TESTING_CUBE_URL }}
           CUBESQL_REWRITE_ENGINE: true
+          CUBESQL_SQL_PUSH_DOWN: true
           CUBESQL_REWRITE_TIMEOUT: 60
         run: cd rust/cubesql && cargo tarpaulin --workspace --no-fail-fast --avoid-cfg-tarpaulin --out Xml
       - name: Upload code coverage

packages/cubejs-api-gateway/src/gateway.ts

@@ -38,7 +38,7 @@ import {
   PreAggJob,
   PreAggJobStatusItem,
   PreAggJobStatusObject,
-  PreAggJobStatusResponse,
+  PreAggJobStatusResponse, SqlApiRequest,
 } from './types/request';
 import {
   CheckAuthInternalOptions,
@@ -1169,7 +1169,7 @@ class ApiGateway {
     return [queryType, normalizedQueries];
   }
 
-  public async sql({ query, context, res }: QueryRequest) {
+  public async sql({ query, context, res, memberToAlias, exportAnnotatedSql }: QueryRequest) {
     const requestStarted = new Date();
 
     try {
@@ -1180,8 +1180,11 @@ class ApiGateway {
 
       const sqlQueries = await Promise.all<any>(
         normalizedQueries.map((normalizedQuery) => this.getCompilerApi(context).getSql(
-          this.coerceForSqlQuery(normalizedQuery, context),
-          { includeDebugInfo: getEnv('devMode') || context.signedWithPlaygroundAuthSecret }
+          this.coerceForSqlQuery({ ...normalizedQuery, memberToAlias }, context),
+          {
+            includeDebugInfo: getEnv('devMode') || context.signedWithPlaygroundAuthSecret,
+            exportAnnotatedSql,
+          }
         ))
       );
 
@@ -1200,6 +1203,30 @@ class ApiGateway {
     }
   }
 
+  public async sqlGenerators({ context, res }: { context: RequestContext, res: ResponseResultFn }) {
+    const requestStarted = new Date();
+
+    try {
+      const compilerApi = this.getCompilerApi(context);
+      const query = {
+        requestId: context.requestId,
+      };
+      const cubeNameToDataSource = await compilerApi.cubeNameToDataSource(query);
+
+      let dataSources = Object.keys(cubeNameToDataSource).map(c => cubeNameToDataSource[c]);
+      dataSources = [...new Set(dataSources)];
+      const dataSourceToSqlGenerator = (await Promise.all(
+        dataSources.map(async dataSource => ({ [dataSource]: (await compilerApi.getSqlGenerator(query, dataSource)).sqlGenerator }))
+      )).reduce((a, b) => ({ ...a, ...b }), {});
+
+      res({ cubeNameToDataSource, dataSourceToSqlGenerator });
+    } catch (e) {
+      this.handleError({
+        e, context, res, requestStarted
+      });
+    }
+  }
+
   protected createSecurityContextExtractor(options?: JWTOptions): SecurityContextExtractorFn {
     if (options?.claimsNamespace) {
       return (ctx: Readonly<RequestContext>) => {
@@ -1609,6 +1636,161 @@ class ApiGateway {
     }
   }
 
+  public async sqlApiLoad(request: SqlApiRequest) {
+    let query: Query | Query[] | null = null;
+    const {
+      context,
+      res,
+      apiType = 'sql',
+    } = request;
+    const requestStarted = new Date();
+
+    try {
+      await this.assertApiScope('data', context.securityContext);
+
+      query = this.parseQueryParam(request.query);
+      let resType: ResultType = ResultType.DEFAULT;
+
+      if (!Array.isArray(query) && query.responseFormat) {
+        resType = query.responseFormat;
+      }
+
+      this.log({
+        type: 'Load Request',
+        query,
+        streaming: request.streaming,
+      }, context);
+
+      const [queryType, normalizedQueries] =
+        await this.getNormalizedQueries(query, context);
+
+      let metaConfigResult = await this
+        .getCompilerApi(context).metaConfig({
+          requestId: context.requestId
+        });
+
+      metaConfigResult = this.filterVisibleItemsInMeta(context, metaConfigResult);
+
+      const sqlQueries = await this
+        .getSqlQueriesInternal(context, normalizedQueries);
+
+      let results;
+
+      let slowQuery = false;
+
+      const streamResponse = async (sqlQuery) => {
+        const q = {
+          ...sqlQuery,
+          query: sqlQuery.query || sqlQuery.sql[0],
+          values: sqlQuery.values || sqlQuery.sql[1],
+          continueWait: true,
+          renewQuery: false,
+          requestId: context.requestId,
+          context,
+          persistent: true,
+          forceNoCache: true,
+        };
+        return {
+          stream: await this.getAdapterApi(context).streamQuery(q),
+        };
+      };
+
+      if (request.sqlQuery) {
+        const finalQuery = {
+          query: request.sqlQuery[0],
+          values: request.sqlQuery[1],
+          continueWait: true,
+          renewQuery: normalizedQueries[0].renewQuery,
+          requestId: context.requestId,
+          context,
+          ...sqlQueries[0],
+          // TODO Can we just pass through data? Ensure hidden members can't be queried
+          aliasNameToMember: null,
+        };
+
+        if (request.streaming) {
+          results = [await streamResponse(finalQuery)];
+        } else {
+          const response = await this
+            .getAdapterApi(context)
+            .executeQuery(finalQuery);
+
+          const annotation = prepareAnnotation(
+            metaConfigResult, normalizedQueries[0]
+          );
+
+          // TODO Can we just pass through data? Ensure hidden members can't be queried
+          results = [{
+            data: response.data,
+            annotation
+          }];
+        }
+      } else {
+        results = await Promise.all(
+          normalizedQueries.map(async (normalizedQuery, index) => {
+            slowQuery = slowQuery ||
+              Boolean(sqlQueries[index].slowQuery);
+
+            const annotation = prepareAnnotation(
+              metaConfigResult, normalizedQuery
+            );
+
+            if (request.streaming) {
+              return streamResponse(sqlQueries[index]);
+            }
+
+            const response = await this.getSqlResponseInternal(
+              context,
+              normalizedQuery,
+              sqlQueries[index],
+            );
+
+            return this.getResultInternal(
+              context,
+              queryType,
+              normalizedQuery,
+              sqlQueries[index],
+              annotation,
+              response,
+              resType,
+            );
+          })
+        );
+      }
+
+      this.log(
+        {
+          type: 'Load Request Success',
+          query,
+          duration: this.duration(requestStarted),
+          apiType,
+          isPlayground: Boolean(
+            context.signedWithPlaygroundAuthSecret
+          ),
+          queries: results.length,
+          queriesWithPreAggregations:
+          results.filter(
+            (r: any) => Object.keys(
+              r.usedPreAggregations || {}
+            ).length
+          ).length,
+          queriesWithData:
+          results.filter((r: any) => r.data?.length).length,
+          dbType: results.map(r => r.dbType),
+        },
+        context,
+      );
+
+      res(request.streaming ? results[0] : {
+        results,
+      });
+    } catch (e) {
+      this.handleError({
+        e, context, query, res, requestStarted
+      });
+    }
+  }
+
   public subscribeQueueEvents({ context, signedWithPlaygroundAuthSecret, connectionId, res }) {
     if (this.enforceSecurityChecks && !signedWithPlaygroundAuthSecret) {
       throw new CubejsHandlerError(

packages/cubejs-api-gateway/src/helpers/transformData.ts

@@ -204,10 +204,6 @@ function getCompactRow(
       getDateRangeValue(timeDimensions)
     );
   } else if (queryType === QueryTypeEnum.BLENDING_QUERY) {
-    console.log(getBlendingResponseKey(timeDimensions));
-    console.log(dbRow[
-      getBlendingResponseKey(timeDimensions)
-    ]);
     row.push(
       dbRow[
         membersToAliasMap[

packages/cubejs-api-gateway/src/sql-server.ts

@@ -45,6 +45,24 @@ export class SQLServer {
     const canSwitchSqlUser: CanSwitchSQLUserFn = options.canSwitchSqlUser
       || this.createDefaultCanSwitchSqlUserFn(options);
 
+    const contextByRequest = async (request, session) => {
+      let userForContext = session.user;
+
+      if (request.meta.changeUser && request.meta.changeUser !== session.user) {
+        const canSwitch = session.superuser || await canSwitchSqlUser(session.user, request.meta.changeUser);
+        if (canSwitch) {
+          userForContext = request.meta.changeUser;
+        } else {
+          throw new Error(
+            `You cannot change security context via __user from ${session.user} to ${request.meta.changeUser}, because it's not allowed.`
+          );
+        }
+      }
+      // @todo Store security context in native for session's user, but not for switching
+      const current = await checkSqlAuth(request, userForContext);
+      return this.contextByNativeReq(request, current.securityContext, request.id);
+    };
+
     this.sqlInterfaceInstance = await registerInterface({
       port: options.sqlPort,
       pgPort: options.pgSqlPort,
@@ -78,22 +96,7 @@ export class SQLServer {
         });
       },
       load: async ({ request, session, query }) => {
-        let userForContext = session.user;
-
-        if (request.meta.changeUser && request.meta.changeUser !== session.user) {
-          const canSwitch = session.superuser || await canSwitchSqlUser(session.user, request.meta.changeUser);
-          if (canSwitch) {
-            userForContext = request.meta.changeUser;
-          } else {
-            throw new Error(
-              `You cannot change security context via __user from ${session.user} to ${request.meta.changeUser}, because it's not allowed.`
-            );
-          }
-        }
-
-        // @todo Store security context in native for session's user, but not for switching
-        const current = await checkSqlAuth(request, userForContext);
-        const context = await this.contextByNativeReq(request, current.securityContext, request.id);
+        const context = await contextByRequest(request, session);
 
         // eslint-disable-next-line no-async-promise-executor
         return new Promise(async (resolve, reject) => {
@@ -112,23 +115,51 @@ export class SQLServer {
           }
         });
       },
-      stream: async ({ request, session, query }) => {
-        let userForContext = session.user;
-
-        if (request.meta.changeUser && request.meta.changeUser !== session.user) {
-          const canSwitch = session.superuser || await canSwitchSqlUser(session.user, request.meta.changeUser);
-          if (canSwitch) {
-            userForContext = request.meta.changeUser;
-          } else {
-            throw new Error(
-              `You cannot change security context via __user from ${session.user} to ${request.meta.changeUser}, because it's not allowed.`
-            );
+      sqlApiLoad: async ({ request, session, query, sqlQuery, streaming }) => {
+        const context = await contextByRequest(request, session);
+
+        // eslint-disable-next-line no-async-promise-executor
+        return new Promise(async (resolve, reject) => {
+          try {
+            await this.apiGateway.sqlApiLoad({
+              query,
+              sqlQuery,
+              streaming,
+              context,
+              res: (message) => {
+                resolve(message);
+              },
+              apiType: 'sql',
+            });
+          } catch (e) {
+            reject(e);
           }
-        }
+        });
+      },
+      sql: async ({ request, session, query, memberToAlias }) => {
+        const context = await contextByRequest(request, session);
 
-        // @todo Store security context in native for session's user, but not for switching
-        const current = await checkSqlAuth(request, userForContext);
-        const context = await this.contextByNativeReq(request, current.securityContext, request.id);
+        // eslint-disable-next-line no-async-promise-executor
+        return new Promise(async (resolve, reject) => {
+          try {
+            await this.apiGateway.sql({
+              query,
+              memberToAlias,
+              exportAnnotatedSql: true,
+              queryType: 'multi',
+              context,
+              res: (message) => {
+                resolve(message);
+              },
+              apiType: 'sql',
+            });
+          } catch (e) {
+            reject(e);
+          }
+        });
+      },
+      stream: async ({ request, session, query }) => {
+        const context = await contextByRequest(request, session);
 
         // eslint-disable-next-line no-async-promise-executor
         return new Promise(async (resolve, reject) => {
@@ -139,6 +170,27 @@ export class SQLServer {
           }
         });
       },
+      sqlGenerators: async (paramsJson: string) => {
+        // TODO get rid of it
+        const { request, session } = JSON.parse(paramsJson);
+        // @todo Store security context in native
+        const { securityContext } = await checkSqlAuth(request, session.user);
+        const context = await this.apiGateway.contextByReq(<any> request, securityContext, request.id);
+
+        // eslint-disable-next-line no-async-promise-executor
+        return new Promise(async (resolve, reject) => {
+          try {
+            await this.apiGateway.sqlGenerators({
+              context,
+              res: (queries) => {
+                resolve(queries);
+              },
+            });
+          } catch (e) {
+            reject(e);
+          }
+        });
+      },
     });
   }