feat: support context_to_roles in Python configuration

Author: bsod90

packages/cubejs-backend-native/js/index.ts

@@ -354,6 +354,7 @@ export interface PyConfiguration {
   checkAuth?: (req: unknown, authorization: string) => Promise<{ 'security_context'?: unknown }>
   queryRewrite?: (query: unknown, ctx: unknown) => Promise<unknown>
   contextToApiScopes?: () => Promise<string[]>
+  contextToRoles?: (ctx: unknown) => Promise<string[]>
 }
 
 function simplifyExpressRequest(req: ExpressRequest) {

packages/cubejs-backend-native/python/cube/src/__init__.py

@@ -74,8 +74,9 @@ class Configuration:
     repository_factory: Callable
     schema_version: Union[str, Callable[[RequestContext], str]]
     semantic_layer_sync: Union[Dict, Callable[[], Dict]]
-    pre_aggregations_schema: Union[Callable[[RequestContext], str]]
+    pre_aggregations_schema: Union[Callable[[RequestContext], str], str]
     orchestrator_options: Union[Dict, Callable[[RequestContext], Dict]]
+    context_to_roles: Callable[[RequestContext], list[str]]
 
     def __init__(self):
         self.web_sockets = None
@@ -123,6 +124,7 @@ def __init__(self):
         self.semantic_layer_sync = None
         self.pre_aggregations_schema = None
         self.orchestrator_options = None
+        self.context_to_roles = None
 
     def __call__(self, func):
         if isinstance(func, str):

packages/cubejs-backend-native/src/python/cube_config.rs

@@ -17,51 +17,52 @@ impl CubeConfigPy {
 
     pub fn get_attrs(&self) -> Vec<&'static str> {
         vec![
-            "web_sockets",
-            "http",
-            "graceful_shutdown",
-            "process_subscriptions_interval",
-            "web_sockets_base_path",
-            "schema_path",
-            "base_path",
-            "dev_server",
+            "allow_js_duplicate_props_in_schema",
             "api_secret",
+            "base_path",
             "cache_and_queue_driver",
-            "allow_js_duplicate_props_in_schema",
+            "compiler_cache_size",
+            "dev_server",
+            "graceful_shutdown",
+            "http",
             "jwt",
+            "live_preview",
+            "max_compiler_cache_keep_alive",
+            "pg_sql_port",
+            "process_subscriptions_interval",
+            "scheduled_refresh_batch_size",
+            "scheduled_refresh_concurrency",
             "scheduled_refresh_timer",
             "scheduled_refresh_timezones",
-            "scheduled_refresh_concurrency",
-            "scheduled_refresh_batch_size",
-            "compiler_cache_size",
-            "update_compiler_cache_keep_alive",
-            "max_compiler_cache_keep_alive",
-            "telemetry",
+            "schema_path",
             "sql_cache",
-            "live_preview",
-            "pg_sql_port",
+            "sql_password",
             "sql_super_user",
             "sql_user",
-            "sql_password",
+            "telemetry",
+            "update_compiler_cache_keep_alive",
+            "web_sockets",
+            "web_sockets_base_path",
             // functions
-            "logger",
+            "can_switch_sql_user",
+            "check_auth",
+            "check_sql_auth",
+            "context_to_api_scopes",
             "context_to_app_id",
             "context_to_orchestrator_id",
+            "context_to_roles",
+            "db_type",
             "driver_factory",
+            "extend_context",
             "external_driver_factory",
-            "db_type",
-            "check_auth",
-            "check_sql_auth",
-            "can_switch_sql_user",
+            "logger",
+            "orchestrator_options",
+            "pre_aggregations_schema",
             "query_rewrite",
-            "extend_context",
-            "scheduled_refresh_contexts",
-            "context_to_api_scopes",
             "repository_factory",
-            "semantic_layer_sync",
+            "scheduled_refresh_contexts",
             "schema_version",
-            "pre_aggregations_schema",
-            "orchestrator_options",
+            "semantic_layer_sync",
         ]
     }
 

packages/cubejs-backend-native/test/config.py

@@ -1,52 +1,61 @@
-from cube import (
-    config,
-    file_repository
-)
+from cube import config, file_repository
 
 config.schema_path = "models"
 config.pg_sql_port = 5555
 config.telemetry = False
 
+
 @config
 def query_rewrite(query, ctx):
-    print('[python] query_rewrite query=', query, ' ctx=', ctx)
+    print("[python] query_rewrite query=", query, " ctx=", ctx)
     return query
 
+
 @config
 async def check_auth(req, authorization):
-    print('[python] check_auth req=', req, ' authorization=', authorization)
+    print("[python] check_auth req=", req, " authorization=", authorization)
     return {
-      'security_context': {
-        'sub': '1234567890',
-        'iat': 1516239022,
-        'user_id': 42
-      },
-      'ignoredField': 'should not be visible'
+        "security_context": {"sub": "1234567890", "iat": 1516239022, "user_id": 42},
+        "ignoredField": "should not be visible",
     }
 
+
 @config
 async def repository_factory(ctx):
-    print('[python] repository_factory ctx=', ctx)
+    print("[python] repository_factory ctx=", ctx)
+
+    return file_repository(ctx["securityContext"]["schemaPath"])
 
-    return file_repository(ctx['securityContext']['schemaPath'])
 
 @config
 async def context_to_api_scopes():
-    print('[python] context_to_api_scopes')
-    return ['meta', 'data', 'jobs']
+    print("[python] context_to_api_scopes")
+    return ["meta", "data", "jobs"]
+
 
 @config
 def schema_version(ctx):
-    print('[python] schema_version', ctx)
+    print("[python] schema_version", ctx)
+
+    return "1"
 
-    return '1'
 
 @config
 def pre_aggregations_schema(ctx):
-    print('[python] pre_aggregations_schema', ctx)
+    print("[python] pre_aggregations_schema", ctx)
+
+    return "schema"
 
-    return 'schema'
 
 @config
 def logger(msg, params):
-    print('[python] logger msg', msg, 'params=', params)
+    print("[python] logger msg", msg, "params=", params)
+
+
+@config
+def context_to_roles(ctx):
+    print("[python] context_to_roles", ctx)
+
+    return [
+        "admin",
+    ]

packages/cubejs-backend-native/test/python.test.ts

@@ -66,6 +66,14 @@ suite('Python Config', () => {
     });
   });
 
+  test('context_to_roles', async () => {
+    if (!config.contextToRoles) {
+      throw new Error('contextToRoles was not defined in config.py');
+    }
+
+    expect(await config.contextToRoles({})).toEqual(['admin']);
+  });
+
   test('context_to_api_scopes', async () => {
     if (!config.contextToApiScopes) {
       throw new Error('contextToApiScopes was not defined in config.py');