chore(server): Support async/promises for ContextAcceptor (#7109)

cube-js · Sep 7, 2023 · 54bb2d9 · 54bb2d9
1 parent 5516f23
commit 54bb2d9
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 13 deletions.
diff --git a/packages/cubejs-api-gateway/src/SubscriptionServer.ts b/packages/cubejs-api-gateway/src/SubscriptionServer.ts
@@ -55,7 +55,7 @@ export class SubscriptionServer {
       if (message.authorization) {
         authContext = { isSubscription: true };
         await this.apiGateway.checkAuthFn(authContext, message.authorization);
-        const acceptanceResult = this.contextAcceptor(authContext);
+        const acceptanceResult = await this.contextAcceptor(authContext);
         if (!acceptanceResult.accepted) {
           this.sendMessage(connectionId, acceptanceResult.rejectMessage);
           return;

diff --git a/packages/cubejs-api-gateway/src/interfaces.ts b/packages/cubejs-api-gateway/src/interfaces.ts
@@ -91,11 +91,13 @@ export type ContextRejectionMiddlewareFn =
    next: ExpressNextFunction,
  ) => void;
 
+type ContextAcceptorResult = { accepted: boolean; rejectMessage?: any };
+
 /**
  * ContextAcceptorFn type that matches the ContextAcceptor.shouldAcceptWs
  * signature from the server-core package
  */
-export type ContextAcceptorFn = (context: RequestContext) => { accepted: boolean; rejectMessage?: any };
+export type ContextAcceptorFn = (context: RequestContext) => Promise<ContextAcceptorResult> | ContextAcceptorResult;
 
 /**
  * Logger middleware.

diff --git a/packages/cubejs-server-core/src/core/server.ts b/packages/cubejs-server-core/src/core/server.ts
@@ -59,16 +59,16 @@ function wrapToFnIfNeeded<T, R>(possibleFn: T | ((a: R) => T)): (a: R) => T {
   return () => possibleFn;
 }
 
-class AcceptAllAcceptor {
-  public shouldAccept(): ContextAcceptanceResult {
+class AcceptAllAcceptor implements ContextAcceptor {
+  public async shouldAccept(): Promise<ContextAcceptanceResult> {
     return { accepted: true };
   }
 
-  public shouldAcceptHttp(): ContextAcceptanceResultHttp {
+  public async shouldAcceptHttp(): Promise<ContextAcceptanceResultHttp> {
     return { accepted: true };
   }
 
-  public shouldAcceptWs(): ContextAcceptanceResultWs {
+  public async shouldAcceptWs(): Promise<ContextAcceptanceResultWs> {
     return { accepted: true };
   }
 }
@@ -475,7 +475,7 @@ export class CubejsServerCore {
 
   protected async contextRejectionMiddleware(req, res, next) {
     if (!this.standalone) {
-      const result = this.contextAcceptor.shouldAcceptHttp(req.context);
+      const result = await this.contextAcceptor.shouldAcceptHttp(req.context);
       if (!result.accepted) {
         res.writeHead(result.rejectStatusCode!, result.rejectHeaders!);
         res.send();
@@ -707,9 +707,17 @@ export class CubejsServerCore {
         error: 'At least one context should be returned by scheduledRefreshContexts'
       });
     }
-    const contexts = allContexts.filter(
-      (context) => this.contextAcceptor.shouldAccept(this.migrateBackgroundContext(context)).accepted
-    );
+
+    const contexts = [];
+
+    for (const allContext of allContexts) {
+      const res = await this.contextAcceptor.shouldAccept(
+        this.migrateBackgroundContext(allContext)
+      );
+      if (res.accepted) {
+        contexts.push(allContext);
+      }
+    }
 
     const batchLimit = pLimit(this.options.scheduledRefreshBatchSize);
     return Promise.all(

diff --git a/packages/cubejs-server-core/src/core/types.ts b/packages/cubejs-server-core/src/core/types.ts
@@ -256,7 +256,7 @@ export type ContextAcceptanceResultWs = ContextAcceptanceResult & {
 };
 
 export interface ContextAcceptor {
-  shouldAccept(context: RequestContext | null): ContextAcceptanceResult;
-  shouldAcceptHttp(context: RequestContext | null): ContextAcceptanceResultHttp;
-  shouldAcceptWs(context: RequestContext | null): ContextAcceptanceResultWs;
+  shouldAccept(context: RequestContext | null): Promise<ContextAcceptanceResult> | ContextAcceptanceResult;
+  shouldAcceptHttp(context: RequestContext | null): Promise<ContextAcceptanceResultHttp> | ContextAcceptanceResultHttp;
+  shouldAcceptWs(context: RequestContext | null): Promise<ContextAcceptanceResultWs> | ContextAcceptanceResultWs;
 }