From 870080732fa1fe1786d2cda3022250771ad9d58a Mon Sep 17 00:00:00 2001
From: Maxim Leonovich <maxim@cube.dev>
Date: Thu, 6 Feb 2025 12:30:17 -0800
Subject: [PATCH] feat(cubejs-server-core): add contextToCubeStoreRouterId
 config option

---
 packages/cubejs-server-core/src/core/optionsValidate.ts | 1 +
 packages/cubejs-server-core/src/core/server.ts          | 4 ++++
 packages/cubejs-server-core/src/core/types.ts           | 2 ++
 3 files changed, 7 insertions(+)

diff --git a/packages/cubejs-server-core/src/core/optionsValidate.ts b/packages/cubejs-server-core/src/core/optionsValidate.ts
index 1cf479adee4bc..0734271731afc 100644
--- a/packages/cubejs-server-core/src/core/optionsValidate.ts
+++ b/packages/cubejs-server-core/src/core/optionsValidate.ts
@@ -74,6 +74,7 @@ const schemaOptions = Joi.object().keys({
   contextToAppId: Joi.func(),
   contextToRoles: Joi.func(),
   contextToOrchestratorId: Joi.func(),
+  contextToCubeStoreRouterId: Joi.func(),
   contextToDataSourceId: Joi.func(),
   contextToApiScopes: Joi.func(),
   repositoryFactory: Joi.func(),
diff --git a/packages/cubejs-server-core/src/core/server.ts b/packages/cubejs-server-core/src/core/server.ts
index c50988ebb0c85..ff576c175fcb4 100644
--- a/packages/cubejs-server-core/src/core/server.ts
+++ b/packages/cubejs-server-core/src/core/server.ts
@@ -59,6 +59,7 @@ import type {
   LoggerFn,
   DriverConfig,
   ScheduledRefreshTimeZonesFn,
+  ContextToCubeStoreRouterIdFn,
 } from './types';
 import {
   ContextToOrchestratorIdFn,
@@ -137,6 +138,8 @@ export class CubejsServerCore {
 
   protected readonly contextToOrchestratorId: ContextToOrchestratorIdFn;
 
+  protected readonly contextToCubeStoreRouterId: ContextToCubeStoreRouterIdFn | null;
+
   protected readonly preAggregationsSchema: PreAggregationsSchemaFn;
 
   protected readonly scheduledRefreshTimeZones: ScheduledRefreshTimeZonesFn;
@@ -216,6 +219,7 @@ export class CubejsServerCore {
     }
 
     this.contextToOrchestratorId = this.options.contextToOrchestratorId || (() => 'STANDALONE');
+    this.contextToCubeStoreRouterId = this.options.contextToCubeStoreRouterId;
 
     // proactively free up old cache values occasionally
     if (this.options.maxCompilerCacheKeepAlive) {
diff --git a/packages/cubejs-server-core/src/core/types.ts b/packages/cubejs-server-core/src/core/types.ts
index d639e2664cf2e..73cdbdd81890b 100644
--- a/packages/cubejs-server-core/src/core/types.ts
+++ b/packages/cubejs-server-core/src/core/types.ts
@@ -123,6 +123,7 @@ export type DatabaseType =
 export type ContextToAppIdFn = (context: RequestContext) => string | Promise<string>;
 export type ContextToRolesFn = (context: RequestContext) => string[] | Promise<string[]>;
 export type ContextToOrchestratorIdFn = (context: RequestContext) => string | Promise<string>;
+export type ContextToCubeStoreRouterIdFn = (context: RequestContext) => string | Promise<string>;
 
 export type OrchestratorOptionsFn = (context: RequestContext) => OrchestratorOptions | Promise<OrchestratorOptions>;
 
@@ -190,6 +191,7 @@ export interface CreateOptions {
   contextToAppId?: ContextToAppIdFn;
   contextToRoles?: ContextToRolesFn;
   contextToOrchestratorId?: ContextToOrchestratorIdFn;
+  contextToCubeStoreRouterId?: ContextToCubeStoreRouterIdFn;
   contextToApiScopes?: ContextToApiScopesFn;
   repositoryFactory?: (context: RequestContext) => SchemaFileRepository;
   checkAuth?: CheckAuthFn;