fix: Refresh Scheduler - empty securityContext in driverFactory (#6316)

Author: ovr

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

@@ -565,12 +565,6 @@ class ApiGateway {
   public async metaExtended({ context, res }: { context: RequestContext, res: ResponseResultFn }) {
     const requestStarted = new Date();
 
-    // TODO: test and remove this function.
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    function visibilityFilter(item) {
-      return getEnv('devMode') || context.signedWithPlaygroundAuthSecret || item.isVisible;
-    }
-
     try {
       await this.assertPermission('meta', context.securityContext);
       const metaConfigExtended = await this.getCompilerApi(context).metaConfigExtended({

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

@@ -38,7 +38,7 @@ type UserBackgroundContext = {
 /**
  * Function that should provides a logic of scheduled returning of
  * the user background context. Used as a part of a main
- * configuration object of the Gateway to provide extendabillity to
+ * configuration object of the Gateway to provide extendability to
  * this logic.
  */
 type ScheduledRefreshContextsFn =

packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts

@@ -47,7 +47,6 @@ export type QueryBody = {
   continueWait?: boolean;
   renewQuery?: boolean;
   requestId?: string;
-  context?: any;
   external?: boolean;
   isJob?: boolean;
   forceNoCache?: boolean;

packages/cubejs-server-core/src/core/OptsHandler.ts

@@ -221,7 +221,7 @@ export class OptsHandler {
           );
         }
         // TODO (buntarb): wrapping this call with assertDriverFactoryResult
-        // change assertions sequince and cause a fail of few tests. Review it.
+        // change assertions sequence and cause a fail of few tests. Review it.
         return this.defaultDriverFactory(ctx);
       } else {
         return this.assertDriverFactoryResult(

packages/cubejs-server-core/src/core/OrchestratorApi.ts

@@ -10,11 +10,11 @@ import {
   QueryBody,
 } from '@cubejs-backend/query-orchestrator';
 
-import { DbTypeAsyncFn, ExternalDbTypeFn, RequestContext } from './types';
+import { DatabaseType, RequestContext } from './types';
 
 export interface OrchestratorApiOptions extends QueryOrchestratorOptions {
-  contextToDbType: DbTypeAsyncFn;
-  contextToExternalDbType: ExternalDbTypeFn;
+  contextToDbType: (dataSource: string) => Promise<DatabaseType>;
+  contextToExternalDbType: () => DatabaseType;
   redisPrefix?: string;
 }
 
@@ -104,34 +104,19 @@ export class OrchestratorApi {
         requestId: query.requestId
       });
 
-      const extractDbType = async (response) => {
-        const dbType = await this.options.contextToDbType({
-          ...query.context,
-          dataSource: response.dataSource,
-        });
-        return dbType;
-      };
-
-      const extractExternalDbType = (response) => (
-        this.options.contextToExternalDbType({
-          ...query.context,
-          dataSource: response.dataSource,
-        })
-      );
-
       if (Array.isArray(data)) {
         const res = await Promise.all(
           data.map(async (item) => ({
             ...item,
-            dbType: await extractDbType(item),
-            extDbType: extractExternalDbType(item),
+            dbType: await this.options.contextToDbType(item.dataSource),
+            extDbType: this.options.contextToExternalDbType(),
           }))
         );
         return res;
       }
 
-      data.dbType = await extractDbType(data);
-      data.extDbType = extractExternalDbType(data);
+      data.dbType = await this.options.contextToDbType(data.dataSource);
+      data.extDbType = this.options.contextToExternalDbType();
 
       return data;
     } catch (err) {

packages/cubejs-server-core/src/core/server.ts

@@ -546,10 +546,11 @@ export class CubejsServerCore {
 
     let externalPreAggregationsDriverPromise: Promise<BaseDriver> | null = null;
 
+    const contextToDbType: DbTypeAsyncFn = this.contextToDbType.bind(this);
     const externalDbType = this.contextToExternalDbType(context);
 
-    // orchestrator options can be empty, if user didnt define it.
-    // so we are adding default and configuring queues concurrencies.
+    // orchestrator options can be empty, if user didn't define it.
+    // so we are adding default and configuring queues concurrency.
     const orchestratorOptions =
       this.optsHandler.getOrchestratorInitializedOptions(
         context,
@@ -638,8 +639,12 @@ export class CubejsServerCore {
             }
           })();
         }),
-        contextToDbType: this.contextToDbType.bind(this),
-        contextToExternalDbType: this.contextToExternalDbType.bind(this),
+        contextToDbType: async (dataSource) => contextToDbType({
+          ...context,
+          dataSource
+        }),
+        // speedup with cache
+        contextToExternalDbType: () => externalDbType,
         redisPrefix: orchestratorId,
         skipExternalCacheAndQueue: externalDbType === 'cubestore',
         cacheAndQueueDriver: this.options.cacheAndQueueDriver,

packages/cubejs-server-core/test/unit/RefreshScheduler.test.ts

@@ -1,6 +1,6 @@
 import R from 'ramda';
 import { BaseDriver } from '@cubejs-backend/query-orchestrator';
-import { SchemaFileRepository } from '@cubejs-backend/shared';
+import { pausePromise, SchemaFileRepository } from '@cubejs-backend/shared';
 import { CubejsServerCore, DatabaseType } from '../../src';
 import { RefreshScheduler } from '../../src/core/RefreshScheduler';
 import { CompilerApi } from '../../src/core/CompilerApi';
@@ -333,34 +333,30 @@ class MockDriver extends BaseDriver {
 
 let testCounter = 1;
 
-const setupScheduler = ({ repository, useOriginalSqlPreAggregations }: { repository: SchemaFileRepository, useOriginalSqlPreAggregations?: boolean }) => {
+const setupScheduler = ({ repository, useOriginalSqlPreAggregations, skipAssertSecurityContext }: { repository: SchemaFileRepository, useOriginalSqlPreAggregations?: boolean, skipAssertSecurityContext?: true }) => {
+  const mockDriver = new MockDriver();
+  const externalDriver = new MockDriver();
+
   const serverCore = new CubejsServerCore({
-    dbType: 'postgres',
     apiSecret: 'foo',
-  });
-  const compilerApi = new CompilerApi(
-    repository,
-    async () => 'postgres',
-    {
-      compileContext: {
-        useOriginalSqlPreAggregations,
-      },
-      logger: (msg, params) => {
-        console.log(msg, params);
-      },
-    }
-  );
+    logger: (msg, params) => console.log(msg, params),
+    driverFactory: async ({ securityContext }) => {
+      expect(typeof securityContext).toEqual('object');
+      if (!skipAssertSecurityContext) {
+        expect(securityContext.hasOwnProperty('tenantId')).toEqual(true);
+      }
 
-  const mockDriver = new MockDriver();
+      return mockDriver;
+    },
+    externalDriverFactory: async ({ securityContext }) => {
+      expect(typeof securityContext).toEqual('object');
+      if (!skipAssertSecurityContext) {
+        expect(securityContext.hasOwnProperty('tenantId')).toEqual(true);
+      }
 
-  const orchestratorApi = new OrchestratorApi(
-    () => mockDriver,
-    (msg, params) => console.log(msg, params),
-    {
-      contextToDbType: async () => 'postgres',
-      contextToExternalDbType(): DatabaseType {
-        return 'cubestore';
-      },
+      return externalDriver;
+    },
+    orchestratorOptions: () => ({
       continueWaitTimeout: 0.1,
       queryCacheOptions: {
         queueOptions: () => ({
@@ -374,14 +370,26 @@ const setupScheduler = ({ repository, useOriginalSqlPreAggregations }: { reposit
         }),
       },
       redisPrefix: `TEST_${testCounter++}`,
+    })
+  });
+
+  const compilerApi = new CompilerApi(
+    repository,
+    async () => 'postgres',
+    {
+      compileContext: {
+        useOriginalSqlPreAggregations,
+      },
+      logger: (msg, params) => {
+        console.log(msg, params);
+      },
     }
   );
 
   jest.spyOn(serverCore, 'getCompilerApi').mockImplementation(() => compilerApi);
-  jest.spyOn(serverCore, 'getOrchestratorApi').mockImplementation(() => <any>orchestratorApi);
 
   const refreshScheduler = new RefreshScheduler(serverCore);
-  return { refreshScheduler, orchestratorApi, mockDriver };
+  return { refreshScheduler, compilerApi, mockDriver };
 };
 
 describe('Refresh Scheduler', () => {
@@ -393,6 +401,11 @@ describe('Refresh Scheduler', () => {
     delete process.env.CUBEJS_DB_QUERY_TIMEOUT;
   });
 
+  afterAll(async () => {
+    // align logs from STDOUT
+    await pausePromise(100);
+  });
+
   test('Round robin pre-aggregation refresh by history priority', async () => {
     process.env.CUBEJS_EXTERNAL_DEFAULT = 'false';
     process.env.CUBEJS_SCHEDULED_REFRESH_DEFAULT = 'true';
@@ -876,6 +889,7 @@ describe('Refresh Scheduler', () => {
     process.env.CUBEJS_SCHEDULED_REFRESH_DEFAULT = 'true';
     const { refreshScheduler } = setupScheduler({
       repository: repositoryWithoutPreAggregations,
+      skipAssertSecurityContext: true,
     });
 
     for (let i = 0; i < 50; i++) {

packages/cubejs-server-core/test/unit/index.test.ts

@@ -167,8 +167,7 @@ describe('index.test', () => {
   test('driverFactory should return driver, failure', async () => {
     const options: CreateOptions = { dbType: () => <any>'mongo', driverFactory: () => <any>null, };
 
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    const [driverFactory, orchestratorOptions] = getCreateOrchestratorOptionsFromServer(options);
+    const [driverFactory, _orchestratorOptions] = getCreateOrchestratorOptionsFromServer(options);
 
     try {
       await driverFactory('default');
@@ -182,8 +181,7 @@ describe('index.test', () => {
   test('externalDriverFactory should return driver, failure', async () => {
     const options: CreateOptions = { dbType: () => <any>'mongo', externalDriverFactory: () => <any>null, };
 
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    const [driverFactory, orchestratorOptions] = getCreateOrchestratorOptionsFromServer(options);
+    const [_driverFactory, orchestratorOptions] = getCreateOrchestratorOptionsFromServer(options);
 
     try {
       await orchestratorOptions.externalDriverFactory();