feat(native): Cube.py - support contextToAppId/contextToOrchestratorId (#6725)

Author: ovr

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

@@ -38,7 +38,8 @@ import {
   PreAggJob,
   PreAggJobStatusItem,
   PreAggJobStatusObject,
-  PreAggJobStatusResponse, SqlApiRequest,
+  PreAggJobStatusResponse,
+  SqlApiRequest,
 } from './types/request';
 import {
   CheckAuthInternalOptions,
@@ -139,8 +140,8 @@ class ApiGateway {
 
   public constructor(
     protected readonly apiSecret: string,
-    protected readonly compilerApi: any,
-    protected readonly adapterApi: any,
+    protected readonly compilerApi: (ctx: RequestContext) => Promise<any>,
+    protected readonly adapterApi: (ctx: RequestContext) => Promise<any>,
     protected readonly logger: any,
     protected readonly options: ApiGatewayOptions,
   ) {
@@ -214,7 +215,7 @@ class ApiGateway {
         },
       ],
       async (req, res) => {
-        const compilerApi = this.getCompilerApi(req.context);
+        const compilerApi = await this.getCompilerApi(req.context);
         let schema = compilerApi.getGraphQLSchema();
         if (!schema) {
           let metaConfig = await compilerApi.metaConfig({
@@ -342,15 +343,13 @@ class ApiGateway {
           }
         },
       ],
-      (req: Request, res: Response) => {
+      async (req: Request, res: Response) => {
         const { transformedQuery, references } = req.body;
-        const canUsePreAggregationForTransformedQuery =
-          this
-            .compilerApi(req.context)
-            .canUsePreAggregationForTransformedQuery(
-              transformedQuery,
-              references,
-            );
+        const compilerApi = await this.getCompilerApi(req.context as RequestContext);
+        const canUsePreAggregationForTransformedQuery = compilerApi.canUsePreAggregationForTransformedQuery(
+          transformedQuery,
+          references,
+        );
         res.json({ canUsePreAggregationForTransformedQuery });
       }
     );
@@ -511,7 +510,8 @@ class ApiGateway {
 
     try {
       await this.assertApiScope('meta', context.securityContext);
-      const metaConfig = await this.getCompilerApi(context).metaConfig({
+      const compilerApi = await this.getCompilerApi(context);
+      const metaConfig = await compilerApi.metaConfig({
         requestId: context.requestId,
       });
       const cubes = this.filterVisibleItemsInMeta(context, metaConfig).map(cube => cube.config);
@@ -531,7 +531,8 @@ class ApiGateway {
 
     try {
       await this.assertApiScope('meta', context.securityContext);
-      const metaConfigExtended = await this.getCompilerApi(context).metaConfigExtended({
+      const compilerApi = await this.getCompilerApi(context);
+      const metaConfigExtended = await compilerApi.metaConfigExtended({
         requestId: context.requestId,
       });
       const { metaConfig, cubeDefinitions } = metaConfigExtended;
@@ -566,7 +567,7 @@ class ApiGateway {
   public async getPreAggregations({ cacheOnly, context, res }: { cacheOnly?: boolean, context: RequestContext, res: ResponseResultFn }) {
     const requestStarted = new Date();
     try {
-      const compilerApi = this.getCompilerApi(context);
+      const compilerApi = await this.getCompilerApi(context);
       const preAggregations = await compilerApi.preAggregations();
 
       const preAggregationPartitions = await this.refreshScheduler()
@@ -600,8 +601,8 @@ class ApiGateway {
         this.parseQueryParam(query),
         { timezones: this.scheduledRefreshTimeZones }
       );
-      const orchestratorApi = this.getAdapterApi(context);
-      const compilerApi = this.getCompilerApi(context);
+      const orchestratorApi = await this.getAdapterApi(context);
+      const compilerApi = await this.getCompilerApi(context);
 
       const preAggregationPartitions = await this.refreshScheduler()
         .preAggregationPartitions(
@@ -661,7 +662,7 @@ class ApiGateway {
       query = normalizeQueryPreAggregationPreview(this.parseQueryParam(query));
       const { preAggregationId, versionEntry, timezone } = query;
 
-      const orchestratorApi = this.getAdapterApi(context);
+      const orchestratorApi = await this.getAdapterApi(context);
 
       const preAggregationPartitions = await this.refreshScheduler()
         .preAggregationPartitions(
@@ -863,7 +864,7 @@ class ApiGateway {
     context: RequestContext,
     selector: PreAggsSelector
   ): Promise<string[]> {
-    const compiler = this.getCompilerApi(context);
+    const compiler = await this.getCompilerApi(context);
     const { timezones } = selector;
     const preaggs = await compiler.preAggregations({
       dataSources: selector.dataSources,
@@ -909,8 +910,8 @@ class ApiGateway {
     const promise: Promise<(PreAggJobStatusItem | undefined)[]> = Promise.all(
       selector.map(async (selected, i) => {
         const ctx = { ...context, ...selected.context };
-        const orchestrator = this.getAdapterApi(ctx);
-        const compiler = this.getCompilerApi(ctx);
+        const orchestrator = await this.getAdapterApi(ctx);
+        const compiler = await this.getCompilerApi(ctx);
         const sel: PreAggsSelector = {
           cubes: [selected.preagg.split('.')[0]],
           preAggregations: [selected.preagg],
@@ -1068,7 +1069,7 @@ class ApiGateway {
   ) {
     const requestStarted = new Date();
     try {
-      const orchestratorApi = this.getAdapterApi(context);
+      const orchestratorApi = await this.getAdapterApi(context);
       res({
         result: await orchestratorApi.getPreAggregationQueueStates()
       });
@@ -1085,7 +1086,7 @@ class ApiGateway {
     const requestStarted = new Date();
     try {
       const { queryKeys, dataSource } = normalizeQueryCancelPreAggregations(this.parseQueryParam(query));
-      const orchestratorApi = this.getAdapterApi(context);
+      const orchestratorApi = await this.getAdapterApi(context);
       res({
         result: await orchestratorApi.cancelPreAggregationQueriesFromQueue(queryKeys, dataSource)
       });
@@ -1179,7 +1180,7 @@ class ApiGateway {
       const [queryType, normalizedQueries] = await this.getNormalizedQueries(query, context);
 
       const sqlQueries = await Promise.all<any>(
-        normalizedQueries.map((normalizedQuery) => this.getCompilerApi(context).getSql(
+        normalizedQueries.map(async (normalizedQuery) => (await this.getCompilerApi(context)).getSql(
           this.coerceForSqlQuery({ ...normalizedQuery, memberToAlias }, context),
           {
             includeDebugInfo: getEnv('devMode') || context.signedWithPlaygroundAuthSecret,
@@ -1207,7 +1208,7 @@ class ApiGateway {
     const requestStarted = new Date();
 
     try {
-      const compilerApi = this.getCompilerApi(context);
+      const compilerApi = await this.getCompilerApi(context);
       const query = {
         requestId: context.requestId,
       };
@@ -1293,7 +1294,7 @@ class ApiGateway {
       const [queryType, normalizedQueries] = await this.getNormalizedQueries(query, context);
 
       const sqlQueries = await Promise.all<any>(
-        normalizedQueries.map((normalizedQuery) => this.getCompilerApi(context).getSql(
+        normalizedQueries.map(async (normalizedQuery) => (await this.getCompilerApi(context)).getSql(
           this.coerceForSqlQuery(normalizedQuery, context),
           {
             includeDebugInfo: getEnv('devMode') || context.signedWithPlaygroundAuthSecret
@@ -1330,7 +1331,7 @@ class ApiGateway {
       normalizedQueries.map(
         async (normalizedQuery, index) => {
           const loadRequestSQLStarted = new Date();
-          const sqlQuery = await this.getCompilerApi(context)
+          const sqlQuery = await (await this.getCompilerApi(context))
             .getSql(
               this.coerceForSqlQuery(normalizedQuery, context)
             );
@@ -1390,8 +1391,7 @@ class ApiGateway {
     }
     const [response, total] = await Promise.all(
       queries.map(async (query) => {
-        const res = await this
-          .getAdapterApi(context)
+        const res = await (await this.getAdapterApi(context))
           .executeQuery(query);
         return res;
       })
@@ -1495,7 +1495,7 @@ class ApiGateway {
         originalQuery: query,
         normalizedQuery: normalizedQueries[0],
         streamingQuery: q,
-        stream: await this.getAdapterApi(context).streamQuery(q),
+        stream: await (await this.getAdapterApi(context)).streamQuery(q),
       };
       return _stream;
     } catch (err: any) {
@@ -1546,10 +1546,10 @@ class ApiGateway {
       const [queryType, normalizedQueries] =
         await this.getNormalizedQueries(query, context);
 
-      let metaConfigResult = await this
-        .getCompilerApi(context).metaConfig({
-          requestId: context.requestId
-        });
+      let metaConfigResult = await (await this
+        .getCompilerApi(context)).metaConfig({
+        requestId: context.requestId
+      });
 
       metaConfigResult = this.filterVisibleItemsInMeta(context, metaConfigResult);
 
@@ -1664,10 +1664,10 @@ class ApiGateway {
       const [queryType, normalizedQueries] =
         await this.getNormalizedQueries(query, context);
 
-      let metaConfigResult = await this
-        .getCompilerApi(context).metaConfig({
-          requestId: context.requestId
-        });
+      const compilerApi = await this.getCompilerApi(context);
+      let metaConfigResult = await compilerApi.metaConfig({
+        requestId: context.requestId
+      });
 
       metaConfigResult = this.filterVisibleItemsInMeta(context, metaConfigResult);
 
@@ -1690,8 +1690,11 @@ class ApiGateway {
           persistent: true,
           forceNoCache: true,
         };
+
+        const adapterApi = await this.getAdapterApi(context);
+
         return {
-          stream: await this.getAdapterApi(context).streamQuery(q),
+          stream: await adapterApi.streamQuery(q),
         };
       };
 
@@ -1711,9 +1714,8 @@ class ApiGateway {
         if (request.streaming) {
           results = [await streamResponse(finalQuery)];
         } else {
-          const response = await this
-            .getAdapterApi(context)
-            .executeQuery(finalQuery);
+          const adapterApi = await this.getAdapterApi(context);
+          const response = await adapterApi.executeQuery(finalQuery);
 
           const annotation = prepareAnnotation(
             metaConfigResult, normalizedQueries[0]
@@ -1791,19 +1793,19 @@ class ApiGateway {
     }
   }
 
-  public subscribeQueueEvents({ context, signedWithPlaygroundAuthSecret, connectionId, res }) {
+  public async subscribeQueueEvents({ context, signedWithPlaygroundAuthSecret, connectionId, res }) {
     if (this.enforceSecurityChecks && !signedWithPlaygroundAuthSecret) {
       throw new CubejsHandlerError(
         403,
         'Forbidden',
         'Only for signed with playground auth secret'
       );
     }
-    return this.getAdapterApi(context).subscribeQueueEvents(connectionId, res);
+    return (await this.getAdapterApi(context)).subscribeQueueEvents(connectionId, res);
   }
 
-  public unSubscribeQueueEvents({ context, connectionId }) {
-    return this.getAdapterApi(context).unSubscribeQueueEvents(connectionId);
+  public async unSubscribeQueueEvents({ context, connectionId }) {
+    return (await this.getAdapterApi(context)).unSubscribeQueueEvents(connectionId);
   }
 
   public async subscribe({
@@ -1866,20 +1868,12 @@ class ApiGateway {
     return query as Query | Query[];
   }
 
-  protected getCompilerApi(context) {
-    if (typeof this.compilerApi === 'function') {
-      return this.compilerApi(context);
-    }
-
-    return this.compilerApi;
+  protected async getCompilerApi(context: RequestContext) {
+    return this.compilerApi(context);
   }
 
-  protected getAdapterApi(context) {
-    if (typeof this.adapterApi === 'function') {
-      return this.adapterApi(context);
-    }
-
-    return this.adapterApi;
+  protected async getAdapterApi(context: RequestContext) {
+    return this.adapterApi(context);
   }
 
   public async contextByReq(req: Request, securityContext, requestId: string): Promise<ExtendedRequestContext> {
@@ -2386,7 +2380,7 @@ class ApiGateway {
     let health: 'HEALTH' | 'DOWN' = 'HEALTH';
 
     if (this.standalone) {
-      const orchestratorApi = await this.adapterApi({});
+      const orchestratorApi = await this.adapterApi({} as any);
 
       try {
         // todo: test other data sources