feat: streaming capabilities (#5995)

Author: buntarb

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

@@ -1,4 +1,5 @@
 /* eslint-disable no-restricted-syntax */
+import * as stream from 'stream';
 import jwt, { Algorithm as JWTAlgorithm } from 'jsonwebtoken';
 import R from 'ramda';
 import bodyParser from 'body-parser';
@@ -86,15 +87,6 @@ import {
   transformPreAggregations,
 } from './helpers/transformMetaExtended';
 
-// const timeoutPromise = (timeout) => (
-//   new Promise((resolve) => (
-//     setTimeout(
-//       () => resolve(null),
-//       timeout,
-//     )
-//   ))
-// );
-
 /**
  * API gateway server class.
  */
@@ -1069,10 +1061,11 @@ class ApiGateway {
   protected async getNormalizedQueries(
     query: Record<string, any> | Record<string, any>[],
     context: RequestContext,
+    persistent = false,
   ): Promise<[QueryType, NormalizedQuery[]]> {
     query = this.parseQueryParam(query);
-    let queryType: QueryType = QueryTypeEnum.REGULAR_QUERY;
 
+    let queryType: QueryType = QueryTypeEnum.REGULAR_QUERY;
     if (!Array.isArray(query)) {
       query = this.compareDateRangeTransformer(query);
       if (Array.isArray(query)) {
@@ -1083,6 +1076,7 @@ class ApiGateway {
     }
 
     const queries = Array.isArray(query) ? query : [query];
+
     const normalizedQueries: NormalizedQuery[] = await Promise.all(
       queries.map(
         async (currentQuery) => validatePostRewrite(
@@ -1094,6 +1088,10 @@ class ApiGateway {
       )
     );
 
+    normalizedQueries.forEach((q) => {
+      this.processQueryLimit(q, persistent);
+    });
+
     if (normalizedQueries.find((currentQuery) => !currentQuery)) {
       throw new Error('queryTransformer returned null query. Please check your queryTransformer implementation');
     }
@@ -1112,6 +1110,29 @@ class ApiGateway {
     return [queryType, normalizedQueries];
   }
 
+  /**
+   * Asserts query limit, sets the default value if neccessary.
+   *
+   * @throw {Error}
+   */
+  public processQueryLimit(query: NormalizedQuery, persistent = false): void {
+    const def = getEnv('dbQueryDefaultLimit') <= getEnv('dbQueryLimit')
+      ? getEnv('dbQueryDefaultLimit')
+      : getEnv('dbQueryLimit');
+
+    if (!persistent) {
+      if (
+        typeof query.limit === 'number' &&
+        query.limit > getEnv('dbQueryLimit')
+      ) {
+        throw new Error('The query limit has been exceeded.');
+      }
+      query.limit = typeof query.limit === 'number'
+        ? query.limit
+        : def;
+    }
+  }
+
   public async sql({ query, context, res }: QueryRequest) {
     const requestStarted = new Date();
 
@@ -1269,7 +1290,6 @@ class ApiGateway {
     context: RequestContext,
     normalizedQuery: NormalizedQuery,
     sqlQuery: any,
-    apiType: string,
   ) {
     const queries = [{
       ...sqlQuery,
@@ -1279,7 +1299,7 @@ class ApiGateway {
       renewQuery: normalizedQuery.renewQuery,
       requestId: context.requestId,
       context,
-      persistent: false, // apiType === 'sql',
+      persistent: false,
     }];
     if (normalizedQuery.total) {
       const normalizedTotal = structuredClone(normalizedQuery);
@@ -1378,6 +1398,50 @@ class ApiGateway {
     };
   }
 
+  /**
+   * Returns stream object which will be used to stream results from
+   * the data source if applicable, returns `null` otherwise.
+   */
+  public async stream(context: RequestContext, query: Query): Promise<null | {
+    originalQuery: Query;
+    normalizedQuery: NormalizedQuery;
+    streamingQuery: unknown;
+    stream: stream.Writable;
+  }> {
+    const requestStarted = new Date();
+    try {
+      this.log({ type: 'Streaming Query', query }, context);
+      const [, normalizedQueries] = await this.getNormalizedQueries(query, context, true);
+      const sqlQuery = (await this.getSqlQueriesInternal(context, normalizedQueries))[0];
+      const q = {
+        ...sqlQuery,
+        query: sqlQuery.sql[0],
+        values: sqlQuery.sql[1],
+        continueWait: true,
+        renewQuery: false,
+        requestId: context.requestId,
+        context,
+        persistent: true,
+        forceNoCache: true,
+      };
+      const _stream = {
+        originalQuery: query,
+        normalizedQuery: normalizedQueries[0],
+        streamingQuery: q,
+        stream: await this.getAdapterApi(context).streamQuery(q),
+      };
+      return _stream;
+    } catch (e) {
+      this.log({
+        type: 'Streaming Error',
+        query,
+        error: (<Error>e).message,
+        duration: this.duration(requestStarted),
+      }, context);
+      return null;
+    }
+  }
+
   /**
    * Data queries APIs (`/load`, `/subscribe`) entry point. Used by
    * `CubejsApi#load` and `CubejsApi#subscribe` methods to fetch the
@@ -1434,7 +1498,6 @@ class ApiGateway {
             context,
             normalizedQuery,
             sqlQueries[index],
-            apiType,
           );
 
           return this.getResultInternal(

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

@@ -112,6 +112,33 @@ 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.`
+            );
+          }
+        }
+
+        // @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 {
+            resolve(await this.apiGateway.stream(context, query));
+          } catch (e) {
+            reject(e);
+          }
+        });
+      },
     });
   }
 

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

@@ -25,7 +25,8 @@ type ApiType =
   'sql' |
   'graphql' |
   'rest' |
-  'ws';
+  'ws' |
+  'stream';
 
 /**
  * Parsed query type data type.

packages/cubejs-api-gateway/test/index.test.ts

@@ -240,6 +240,7 @@ describe('API Gateway', () => {
               timezone: 'UTC',
               order: [],
               filters: [],
+              limit: 10000,
               dimensions: [],
               timeDimensions: [],
               queryType: 'regularQuery'
@@ -251,6 +252,7 @@ describe('API Gateway', () => {
             timezone: 'UTC',
             order: [],
             filters: [],
+            limit: 10000,
             dimensions: [],
             timeDimensions: [],
             queryType: 'regularQuery'

packages/cubejs-backend-native/Cargo.lock

@@ -1,5 +1,6 @@
 import fs from 'fs';
 import path from 'path';
+// import { getEnv } from '@cubejs-backend/shared';
 
 export interface BaseMeta {
     // postgres or mysql
@@ -53,6 +54,7 @@ export type SQLInterfaceOptions = {
     checkAuth: (payload: CheckAuthPayload) => CheckAuthResponse | Promise<CheckAuthResponse>,
     load: (payload: LoadPayload) => unknown | Promise<unknown>,
     meta: (payload: MetaPayload) => unknown | Promise<unknown>,
+    stream: (payload: LoadPayload) => unknown | Promise<unknown>,
 };
 
 function loadNative() {
@@ -106,6 +108,43 @@ function wrapNativeFunctionWithChannelCallback(
     };
 };
 
+// TODO: Refactor - define classes
+function wrapNativeFunctionWithStream(
+    fn: (extra: any) => unknown | Promise<unknown>
+) {
+    return async (extra: any, writer: any) => {
+        let streamResponse: any;
+        try {
+            streamResponse = await fn(JSON.parse(extra));
+            let chunk: object[] = [];
+            streamResponse.stream.on('data', (c: object) => {
+                chunk.push(c);
+                if (chunk.length >= 10000) {
+                    if (!writer.chunk(JSON.stringify(chunk))) {
+                        streamResponse.stream.removeAllListeners();
+                    }
+                    chunk = [];
+                }
+            });
+            streamResponse.stream.on('close', () => {
+                if (chunk.length > 0) {
+                    writer.chunk(JSON.stringify(chunk));
+                }
+                writer.end("");
+            });
+
+            streamResponse.stream.on('error', (err: any) => {
+                writer.reject(err.message || "Unknown JS exception");
+            });
+        } catch (e: any) {
+            if (!!streamResponse && !!streamResponse.stream) {
+                streamResponse.stream.destroy(e);
+            }
+            writer.reject(e.message || "Unknown JS exception");
+        }
+    };
+};
+
 type LogLevel = 'error' | 'warn' | 'info' | 'debug' | 'trace';
 
 export const setupLogger = (logger: (extra: any) => unknown, logLevel: LogLevel): void => {
@@ -132,12 +171,17 @@ export const registerInterface = async (options: SQLInterfaceOptions): Promise<S
         throw new Error('options.meta must be a function');
     }
 
+    if (typeof options.stream != 'function') {
+        throw new Error('options.stream must be a function');
+    }
+
     const native = loadNative();
     return native.registerInterface({
         ...options,
         checkAuth: wrapNativeFunctionWithChannelCallback(options.checkAuth),
         load: wrapNativeFunctionWithChannelCallback(options.load),
         meta: wrapNativeFunctionWithChannelCallback(options.meta),
+        stream: wrapNativeFunctionWithStream(options.stream),
     });
 };
 

packages/cubejs-backend-native/js/index.ts

@@ -39,6 +39,7 @@
   },
   "dependencies": {
     "@cubejs-backend/cubesql": "^0.31.40",
+    "@cubejs-backend/shared": "^0.31.4",
     "@mapbox/node-pre-gyp": "^1"
   },
   "binary": {