feat: Hooks for the context rejection mechanism (#5901)

* feat: Hooks for the context rejection mechanism

* feat: add context acceptor mechanism to the SubscriptionServer

Author: bsod90

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

@@ -3,7 +3,7 @@ import { v4 as uuidv4 } from 'uuid';
 import { UserError } from './UserError';
 import type { ApiGateway } from './gateway';
 import type { LocalSubscriptionStore } from './LocalSubscriptionStore';
-import { ExtendedRequestContext } from './interfaces';
+import { ExtendedRequestContext, ContextAcceptorFn } from './interfaces';
 
 const methodParams: Record<string, string[]> = {
   load: ['query', 'queryType'],
@@ -26,6 +26,7 @@ export class SubscriptionServer {
     protected readonly apiGateway: ApiGateway,
     protected readonly sendMessage: WebSocketSendMessageFn,
     protected readonly subscriptionStore: LocalSubscriptionStore,
+    protected readonly contextAcceptor: ContextAcceptorFn,
   ) {
   }
 
@@ -54,6 +55,11 @@ export class SubscriptionServer {
       if (message.authorization) {
         authContext = { isSubscription: true };
         await this.apiGateway.checkAuthFn(authContext, message.authorization);
+        const acceptanceResult = this.contextAcceptor(authContext);
+        if (!acceptanceResult.accepted) {
+          this.sendMessage(connectionId, acceptanceResult.rejectMessage);
+          return;
+        }
         await this.subscriptionStore.setAuthContext(connectionId, authContext);
         this.sendMessage(connectionId, { handshake: true });
         return;

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

@@ -54,6 +54,8 @@ import {
 import {
   CheckAuthMiddlewareFn,
   RequestLoggerMiddlewareFn,
+  ContextRejectionMiddlewareFn,
+  ContextAcceptorFn,
 } from './interfaces';
 import { getRequestIdFromRequest, requestParser } from './requestParser';
 import { UserError } from './UserError';
@@ -126,6 +128,10 @@ class ApiGateway {
   protected readonly requestLoggerMiddleware: RequestLoggerMiddlewareFn;
 
   protected readonly securityContextExtractor: SecurityContextExtractorFn;
+  
+  protected readonly contextRejectionMiddleware: ContextRejectionMiddlewareFn;
+
+  protected readonly wsContextAcceptor: ContextAcceptorFn;
 
   protected readonly releaseListeners: (() => any)[] = [];
 
@@ -158,12 +164,15 @@ class ApiGateway {
       : this.checkAuth;
     this.securityContextExtractor = this.createSecurityContextExtractor(options.jwt);
     this.requestLoggerMiddleware = options.requestLoggerMiddleware || this.requestLogger;
+    this.contextRejectionMiddleware = options.contextRejectionMiddleware || (async (req, res, next) => next());
+    this.wsContextAcceptor = options.wsContextAcceptor || (() => ({ accepted: true }));
   }
 
   public initApp(app: ExpressApplication) {
     const userMiddlewares: RequestHandler[] = [
       this.checkAuthMiddleware,
       this.requestContextMiddleware,
+      this.contextRejectionMiddleware,
       this.logNetworkUsage,
       this.requestLoggerMiddleware
     ];
@@ -279,6 +288,7 @@ class ApiGateway {
       const systemMiddlewares: RequestHandler[] = [
         this.checkAuthSystemMiddleware,
         this.requestContextMiddleware,
+        this.contextRejectionMiddleware,
         this.requestLoggerMiddleware
       ];
 
@@ -386,7 +396,7 @@ class ApiGateway {
   }
 
   public initSubscriptionServer(sendMessage: WebSocketSendMessageFn) {
-    return new SubscriptionServer(this, sendMessage, this.subscriptionStore);
+    return new SubscriptionServer(this, sendMessage, this.subscriptionStore, this.wsContextAcceptor);
   }
 
   protected duration(requestStarted) {

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

@@ -75,6 +75,22 @@ export type CheckAuthMiddlewareFn =
    next: ExpressNextFunction,
  ) => void;
 
+/**
+ * Context rejection middleware.
+ */
+export type ContextRejectionMiddlewareFn =
+ (
+   req: Request,
+   res: ExpressResponse,
+   next: ExpressNextFunction,
+ ) => void;
+
+/**
+ * ContextAcceptorFn type that matches the ContextAcceptor.shouldAcceptWs
+ * signature from the server-core package
+ */
+export type ContextAcceptorFn = (context: RequestContext) => { accepted: boolean; rejectMessage?: any };
+
 /**
  * Logger middleware.
  * @deprecated

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

@@ -16,6 +16,8 @@ import {
 import {
   CheckAuthMiddlewareFn,
   RequestLoggerMiddlewareFn,
+  ContextRejectionMiddlewareFn,
+  ContextAcceptorFn,
 } from '../interfaces';
 
 type UserBackgroundContext = {
@@ -59,6 +61,8 @@ interface ApiGatewayOptions {
   enforceSecurityChecks?: boolean;
   playgroundAuthSecret?: string;
   serverCoreVersion?: string;
+  contextRejectionMiddleware?: ContextRejectionMiddlewareFn;
+  wsContextAcceptor?: ContextAcceptorFn;
   checkAuth?: CheckAuthFn;
   /**
    * @deprecated Use checkAuth property instead.

packages/cubejs-backend-shared/src/env.ts

@@ -190,6 +190,9 @@ const variables: Record<string, (...args: any) => any> = {
   maxPartitionsPerCube: () => get('CUBEJS_MAX_PARTITIONS_PER_CUBE')
     .default('10000')
     .asInt(),
+  scheduledRefreshBatchSize: () => get('CUBEJS_SCHEDULED_REFRESH_BATCH_SIZE')
+    .default('1')
+    .asInt(),
 
   /** ****************************************************************
    * Common db options                                               *

packages/cubejs-server-core/package.json

@@ -48,6 +48,7 @@
     "lru-cache": "^5.1.1",
     "moment": "^2.29.1",
     "node-fetch": "^2.6.0",
+    "p-limit": "^3.1.0",
     "promise-timeout": "^1.3.0",
     "ramda": "^0.27.0",
     "semver": "^6.3.0",

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

@@ -464,6 +464,7 @@ export class OptsHandler {
       dashboardAppPort: 3000,
       scheduledRefreshConcurrency:
         parseInt(process.env.CUBEJS_SCHEDULED_REFRESH_CONCURRENCY, 10),
+      scheduledRefreshBatchSize: getEnv('scheduledRefreshBatchSize'),
       preAggregationsSchema:
         getEnv('preAggregationsSchema') ||
         (this.isDevMode()

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

@@ -84,6 +84,7 @@ const schemaOptions = Joi.object().keys({
   scheduledRefreshTimeZones: Joi.array().items(Joi.string()),
   scheduledRefreshContexts: Joi.func(),
   scheduledRefreshConcurrency: Joi.number().min(1).integer(),
+  scheduledRefreshBatchSize: Joi.number().min(1).integer(),
   // Compiler cache
   compilerCacheSize: Joi.number().min(0).integer(),
   updateCompilerCacheKeepAlive: Joi.boolean(),

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

@@ -3,6 +3,7 @@ import crypto from 'crypto';
 import fs from 'fs-extra';
 import LRUCache from 'lru-cache';
 import isDocker from 'is-docker';
+import pLimit from 'p-limit';
 
 import { ApiGateway, UserBackgroundContext } from '@cubejs-backend/api-gateway';
 import {
@@ -11,7 +12,7 @@ import {
   getEnv, assertDataSource, getRealType, internalExceptions, track,
 } from '@cubejs-backend/shared';
 
-import type { Application as ExpressApplication } from 'express';
+import type { Application as ExpressApplication, NextFunction } from 'express';
 
 import { BaseDriver, DriverFactoryByDataSource } from '@cubejs-backend/query-orchestrator';
 import { FileRepository, SchemaFileRepository } from './FileRepository';
@@ -47,7 +48,7 @@ import type {
   LoggerFn,
   DriverConfig,
 } from './types';
-import { ContextToOrchestratorIdFn } from './types';
+import { ContextToOrchestratorIdFn, ContextAcceptanceResult, ContextAcceptanceResultHttp, ContextAcceptanceResultWs, ContextAcceptor } from './types';
 
 const { version } = require('../../../package.json');
 
@@ -59,6 +60,20 @@ function wrapToFnIfNeeded<T, R>(possibleFn: T | ((a: R) => T)): (a: R) => T {
   return () => possibleFn;
 }
 
+class AcceptAllAcceptor {
+  public shouldAccept(): ContextAcceptanceResult {
+    return { accepted: true };
+  }
+
+  public shouldAcceptHttp(): ContextAcceptanceResultHttp {
+    return { accepted: true };
+  }
+
+  public shouldAcceptWs(): ContextAcceptanceResultWs {
+    return { accepted: true };
+  }
+}
+
 export class CubejsServerCore {
   /**
    * Returns core version based on package.json.
@@ -67,13 +82,6 @@ export class CubejsServerCore {
     return version;
   }
 
-  /**
-   * Create an instance of the core.
-   */
-  public static create(options?: CreateOptions, systemOptions?: SystemOptions) {
-    return new CubejsServerCore(options, systemOptions);
-  }
-
   /**
    * Resolve driver module name by db type.
    */
@@ -142,6 +150,8 @@ export class CubejsServerCore {
 
   public coreServerVersion: string | null = null;
 
+  private contextAcceptor: ContextAcceptor;
+
   /**
    * Class constructor.
    */
@@ -179,6 +189,8 @@ export class CubejsServerCore {
       this.standalone = false;
     }
 
+    this.contextAcceptor = this.createContextAcceptor();
+
     if (this.options.contextToDataSourceId) {
       throw new Error('contextToDataSourceId has been deprecated and removed. Use contextToOrchestratorId instead.');
     }
@@ -309,6 +321,10 @@ export class CubejsServerCore {
     }
   }
 
+  protected createContextAcceptor(): ContextAcceptor {
+    return new AcceptAllAcceptor();
+  }
+
   /**
    * Determines whether current instance is ready to process queries.
    */
@@ -419,7 +435,7 @@ export class CubejsServerCore {
       return this.apiGatewayInstance;
     }
 
-    return this.apiGatewayInstance = new ApiGateway(
+    return (this.apiGatewayInstance = new ApiGateway(
       this.options.apiSecret,
       this.getCompilerApi.bind(this),
       this.getOrchestratorApi.bind(this),
@@ -429,17 +445,34 @@ export class CubejsServerCore {
         dataSourceStorage: this.orchestratorStorage,
         basePath: this.options.basePath,
         checkAuthMiddleware: this.options.checkAuthMiddleware,
+        contextRejectionMiddleware: this.contextRejectionMiddleware.bind(this),
+        wsContextAcceptor: this.contextAcceptor.shouldAcceptWs,
         checkAuth: this.options.checkAuth,
-        queryRewrite: this.options.queryRewrite || this.options.queryTransformer,
+        queryRewrite:
+          this.options.queryRewrite || this.options.queryTransformer,
         extendContext: this.options.extendContext,
         playgroundAuthSecret: getEnv('playgroundAuthSecret'),
         jwt: this.options.jwt,
         refreshScheduler: () => new RefreshScheduler(this),
         scheduledRefreshContexts: this.options.scheduledRefreshContexts,
         scheduledRefreshTimeZones: this.options.scheduledRefreshTimeZones,
-        serverCoreVersion: this.coreServerVersion
+        serverCoreVersion: this.coreServerVersion,
       }
-    );
+    ));
+  }
+
+  protected async contextRejectionMiddleware(req, res, next) {
+    if (!this.standalone) {
+      const result = this.contextAcceptor.shouldAcceptHttp(req.context);
+      if (!result.accepted) {
+        res.writeHead(result.rejectStatusCode!, result.rejectHeaders!);
+        res.send();
+        return;
+      }
+    }
+    if (next) {
+      next();
+    }
   }
 
   public getCompilerApi(context: RequestContext) {
@@ -648,22 +681,33 @@ export class CubejsServerCore {
    * @internal Please dont use this method directly, use refreshTimer
    */
   public handleScheduledRefreshInterval = async (options) => {
-    const contexts = await this.options.scheduledRefreshContexts();
+    const contexts = (await this.options.scheduledRefreshContexts()).filter(
+      (context) => this.contextAcceptor.shouldAccept(this.migrateBackgroundContext(context)).accepted
+    );
     if (contexts.length < 1) {
       this.logger('Refresh Scheduler Error', {
         error: 'At least one context should be returned by scheduledRefreshContexts'
       });
     }
 
-    return Promise.all(contexts.map(async context => {
-      const queryingOptions: any = { ...options, concurrency: this.options.scheduledRefreshConcurrency };
+    const batchLimit = pLimit(this.options.scheduledRefreshBatchSize);
+    return Promise.all(
+      contexts
+        .map((context) => async () => {
+          const queryingOptions: any = {
+            ...options,
+            concurrency: this.options.scheduledRefreshConcurrency,
+          };
 
-      if (this.options.scheduledRefreshTimeZones) {
-        queryingOptions.timezones = this.options.scheduledRefreshTimeZones;
-      }
+          if (this.options.scheduledRefreshTimeZones) {
+            queryingOptions.timezones = this.options.scheduledRefreshTimeZones;
+          }
 
-      return this.runScheduledRefresh(context, queryingOptions);
-    }));
+          return this.runScheduledRefresh(context, queryingOptions);
+        })
+        // Limit the number of refresh contexts we process per iteration
+        .map(batchLimit)
+    );
   };
 
   protected getRefreshScheduler() {

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

@@ -187,6 +187,7 @@ export interface CreateOptions {
   scheduledRefreshTimeZones?: string[];
   scheduledRefreshContexts?: () => Promise<UserBackgroundContext[]>;
   scheduledRefreshConcurrency?: number;
+  scheduledRefreshBatchSize?: number;
   compilerCacheSize?: number;
   maxCompilerCacheKeepAlive?: number;
   updateCompilerCacheKeepAlive?: boolean;
@@ -228,3 +229,23 @@ export type ServerCoreInitializedOptions = Required<
 export type SystemOptions = {
   isCubeConfigEmpty: boolean;
 };
+
+// Types to support the ContextAcceptance mechanism
+export type ContextAcceptanceResult = {
+  accepted: boolean;
+};
+
+export type ContextAcceptanceResultHttp = ContextAcceptanceResult & {
+  rejectHeaders?: { [key: string]: string };
+  rejectStatusCode?: number;
+};
+
+export type ContextAcceptanceResultWs = ContextAcceptanceResult & {
+  rejectMessage?: any;
+};
+
+export interface ContextAcceptor {
+  shouldAccept(context: RequestContext | null): ContextAcceptanceResult;
+  shouldAcceptHttp(context: RequestContext | null): ContextAcceptanceResultHttp;
+  shouldAcceptWs(context: RequestContext | null): ContextAcceptanceResultWs;
+}