chore(native): Support python execution in a separate scope (#6899)

ovr · web-flow · commit 414c9f2d4d68 · 2023-07-20T19:06:57.000+03:00
diff --git a/.github/workflows/rust-cubesql.yml b/.github/workflows/rust-cubesql.yml
@@ -31,12 +31,12 @@ jobs:
           toolchain: nightly-2022-03-22
           override: true
           components: rustfmt, clippy
-      - uses: Swatinem/rust-cache@v1
+      - uses: Swatinem/rust-cache@v2
         with:
-          working-directory: ./rust/cubesql
+          workspaces: ./rust/cubesql -> target
           # default key
           key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
-          sharedKey: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
+          shared-key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
       - name: Lint CubeSQL
         run: cd rust/cubesql/cubesql && cargo fmt --all -- --check
       - name: Lint Native
@@ -64,12 +64,12 @@ jobs:
           toolchain: nightly-2022-03-22
           override: true
           components: rustfmt
-      - uses: Swatinem/rust-cache@v1
+      - uses: Swatinem/rust-cache@v2
         with:
-          working-directory: ./rust/cubesql
+          workspaces: ./rust/cubesql -> target
           # default key
           key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
-          sharedKey: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
+          shared-key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
       - name: Install tarpaulin@0.20.1
         uses: baptiste0928/cargo-install@v1
         with:
@@ -107,12 +107,12 @@ jobs:
           toolchain: nightly-2022-03-22
           override: true
           components: rustfmt
-      - uses: Swatinem/rust-cache@v1
+      - uses: Swatinem/rust-cache@v2
         with:
-          working-directory: ./rust/cubesql
+          workspaces: ./rust/cubesql -> target
           # default key
           key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
-          sharedKey: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
+          shared-key: cubesql-${{ runner.OS }}-x86_64-unknown-linux-gnu-16
       - name: Unit tests (Legacy Engine)
         env:
           CUBESQL_TESTING_CUBE_TOKEN: ${{ secrets.CUBESQL_TESTING_CUBE_TOKEN }}
@@ -151,11 +151,11 @@ jobs:
           override: true
           components: rustfmt
           target: ${{ matrix.target }}
-      - uses: Swatinem/rust-cache@v1
+      - uses: Swatinem/rust-cache@v2
         with:
-          working-directory: ./rust/cubesql
+          workspaces: ./rust/cubesql -> target
           key: cubesql-${{ runner.OS }}-${{ matrix.target }}-${{ matrix.node-version }}
-          sharedKey: cubesql-${{ runner.OS }}-${{ matrix.target }}-${{ matrix.node-version }}
+          shared-key: cubesql-${{ runner.OS }}-${{ matrix.target }}-${{ matrix.node-version }}
       - name: Install Node.js ${{ matrix.node-version }}
         uses: actions/setup-node@v3
         with:
@@ -190,24 +190,28 @@ jobs:
         if: (matrix.python-version == 'fallback')
         env:
           CARGO_BUILD_TARGET: ${{ matrix.target }}
-        run: cd packages/cubejs-backend-native && yarn run native:build-debug
+        working-directory: ./packages/cubejs-backend-native
+        run: yarn run native:build-debug
       - name: Build native (with Python)
         if: (matrix.python-version != 'fallback')
         env:
           PYO3_PYTHON: python${{ matrix.python-version }}
           CARGO_BUILD_TARGET: ${{ matrix.target }}
-        run: cd packages/cubejs-backend-native && yarn run native:build-debug-python
+        working-directory: ./packages/cubejs-backend-native
+        run: yarn run native:build-debug-python
       - name: Test native (GNU only)
         if: (matrix.target == 'x86_64-unknown-linux-gnu')
         env:
           CUBESQL_STREAM_MODE: true
           CUBEJS_NATIVE_INTERNAL_DEBUG: true
-        run: cd packages/cubejs-backend-native && yarn run test:unit
+        working-directory: ./packages/cubejs-backend-native
+        run: yarn run test:unit
       - name: Run E2E Smoke testing over whole Cube (GNU only)
         if: (matrix.target == 'x86_64-unknown-linux-gnu')
         env:
           CUBEJS_NATIVE_INTERNAL_DEBUG: true
-        run: cd packages/cubejs-testing && yarn smoke:cubesql
+        working-directory: ./packages/cubejs-testing
+        run: yarn smoke:cubesql
 
   native_macos:
     needs: [lint]
diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts
@@ -248,13 +248,13 @@ export interface PyConfiguration {
     contextToApiScopes?: () => Promise<string[]>
 }
 
-export const pythonLoadConfig = async (context: string, options: { file: string }): Promise<PyConfiguration> => {
+export const pythonLoadConfig = async (content: string, options: { file: string }): Promise<PyConfiguration> => {
   if (isFallbackBuild()) {
     throw new Error('Python is not supported in fallback build');
   }
 
   const native = loadNative();
-  const config = await native.pythonLoadConfig(context, options);
+  const config = await native.pythonLoadConfig(content, options);
 
   if (config.checkAuth) {
     const nativeCheckAuth = config.checkAuth;
diff --git a/packages/cubejs-backend-native/src/python/entry.rs b/packages/cubejs-backend-native/src/python/entry.rs
@@ -5,7 +5,7 @@ use crate::python::runtime::py_runtime_init;
 use crate::python::template;
 use neon::prelude::*;
 use pyo3::prelude::*;
-use pyo3::types::{PyFunction, PyList, PyString, PyTuple};
+use pyo3::types::{PyDict, PyFunction, PyList, PyString, PyTuple};
 
 fn python_load_config(mut cx: FunctionContext) -> JsResult<JsPromise> {
     let config_file_content = cx.argument::<JsString>(0)?.value(&mut cx);
@@ -23,7 +23,12 @@ fn python_load_config(mut cx: FunctionContext) -> JsResult<JsPromise> {
         PyModule::from_code(py, cube_conf_code, "__init__.py", "cube.conf")?;
 
         let config_module = PyModule::from_code(py, &config_file_content, "config.py", "")?;
-        let settings_py = config_module.getattr("settings")?;
+        let settings_py = if config_module.hasattr("__execution_context_locals")? {
+            let execution_context_locals = config_module.getattr("__execution_context_locals")?;
+            execution_context_locals.get_item("settings")?
+        } else {
+            config_module.getattr("settings")?
+        };
 
         let mut cube_conf = CubeConfigPy::new();
 
@@ -61,25 +66,43 @@ fn python_load_model(mut cx: FunctionContext) -> JsResult<JsPromise> {
         PyModule::from_code(py, &cube_code, "__init__.py", "cube")?;
 
         let model_module = PyModule::from_code(py, &model_content, &model_file_name, "")?;
-
-        let inspect_module = py.import("inspect")?;
-        let args = (model_module, inspect_module.getattr("isfunction")?);
-        let functions_with_names = inspect_module
-            .call_method1("getmembers", args)?
-            .downcast::<PyList>()?;
         let mut collected_functions = CLReprObject::new();
 
-        for function_details in functions_with_names.iter() {
-            let function_details = function_details.downcast::<PyTuple>()?;
-            let fun_name = function_details.get_item(0)?.downcast::<PyString>()?;
-            let fun = function_details.get_item(1)?.downcast::<PyFunction>()?;
-
-            let has_attr = fun.hasattr("cube_context_func")?;
-            if has_attr {
-                let fun: Py<PyFunction> = fun.into();
-                collected_functions.insert(fun_name.to_string(), CLRepr::PyExternalFunction(fun));
+        if model_module.hasattr("__execution_context_locals")? {
+            let execution_context_locals = model_module
+                .getattr("__execution_context_locals")?
+                .downcast::<PyDict>()?;
+
+            for (local_key, local_value) in execution_context_locals.iter() {
+                if local_value.is_instance_of::<PyFunction>()? {
+                    let has_attr = local_value.hasattr("cube_context_func")?;
+                    if has_attr {
+                        let fun: Py<PyFunction> = local_value.downcast::<PyFunction>()?.into();
+                        collected_functions
+                            .insert(local_key.to_string(), CLRepr::PyExternalFunction(fun));
+                    }
+                }
             }
-        }
+        } else {
+            let inspect_module = py.import("inspect")?;
+            let args = (model_module, inspect_module.getattr("isfunction")?);
+            let functions_with_names = inspect_module
+                .call_method1("getmembers", args)?
+                .downcast::<PyList>()?;
+
+            for function_details in functions_with_names.iter() {
+                let function_details = function_details.downcast::<PyTuple>()?;
+                let fun_name = function_details.get_item(0)?.downcast::<PyString>()?;
+                let fun = function_details.get_item(1)?.downcast::<PyFunction>()?;
+
+                let has_attr = fun.hasattr("cube_context_func")?;
+                if has_attr {
+                    let fun: Py<PyFunction> = fun.into();
+                    collected_functions
+                        .insert(fun_name.to_string(), CLRepr::PyExternalFunction(fun));
+                }
+            }
+        };
 
         Ok(CubePythonModel::new(collected_functions))
     });
diff --git a/packages/cubejs-backend-native/test/jinja.test.ts b/packages/cubejs-backend-native/test/jinja.test.ts
@@ -5,6 +5,9 @@ import * as native from '../js';
 import type { JinjaEngine } from '../js';
 
 const suite = native.isFallbackBuild() ? xdescribe : describe;
+// TODO(ovr): Find what is going wrong with parallel tests & python on Linux
+const darwinSuite = process.platform === 'darwin' && !native.isFallbackBuild() ? describe : xdescribe;
+
 const nativeInstance = new native.NativeInstance();
 
 function loadTemplateFile(engine: native.JinjaEngine, fileName: string): void {
@@ -13,10 +16,10 @@ function loadTemplateFile(engine: native.JinjaEngine, fileName: string): void {
   engine.loadTemplate(fileName, content);
 }
 
-async function loadPythonCtxFromUtils() {
-  const content = fs.readFileSync(path.join(process.cwd(), 'test', 'templates', 'utils.py'), 'utf8');
+async function loadPythonCtxFromUtils(fileName: string) {
+  const content = fs.readFileSync(path.join(process.cwd(), 'test', 'templates', fileName), 'utf8');
   return nativeInstance.loadPythonContext(
-    'utils.py',
+    fileName,
     content
   );
 }
@@ -31,7 +34,7 @@ function testTemplateBySnapshot(engine: JinjaEngine, templateName: string, ctx:
 
 function testTemplateWithPythonCtxBySnapshot(engine: JinjaEngine, templateName: string, ctx: unknown) {
   test(`render ${templateName}`, async () => {
-    const actual = engine.renderTemplate(templateName, ctx, await loadPythonCtxFromUtils());
+    const actual = engine.renderTemplate(templateName, ctx, await loadPythonCtxFromUtils('utils.py'));
 
     expect(actual).toMatchSnapshot(templateName);
   });
@@ -49,9 +52,24 @@ function testLoadBrokenTemplateBySnapshot(engine: JinjaEngine, templateName: str
   });
 }
 
-suite('Python models', () => {
+suite('Python model', () => {
   it('load utils.py', async () => {
-    const pythonModule = await loadPythonCtxFromUtils();
+    const pythonModule = await loadPythonCtxFromUtils('utils.py');
+
+    expect(pythonModule).toEqual({
+      load_data: expect.any(Object),
+      load_data_sync: expect.any(Object),
+      arg_bool: expect.any(Object),
+      arg_sum_integers: expect.any(Object),
+      arg_str: expect.any(Object),
+      arg_null: expect.any(Object),
+    });
+  });
+});
+
+darwinSuite('Scope Python model', () => {
+  it('load scoped-utils.py', async () => {
+    const pythonModule = await loadPythonCtxFromUtils('scoped-utils.py');
 
     expect(pythonModule).toEqual({
       load_data: expect.any(Object),
diff --git a/packages/cubejs-backend-native/test/python.test.ts b/packages/cubejs-backend-native/test/python.test.ts
@@ -5,30 +5,33 @@ import * as native from '../js';
 import { PyConfiguration } from '../js';
 
 const suite = native.isFallbackBuild() ? xdescribe : describe;
+// TODO(ovr): Find what is going wrong with parallel tests & python on Linux
+const darwinSuite = process.platform === 'darwin' && !native.isFallbackBuild() ? describe : xdescribe;
+
+async function loadConfigurationFile(file: string) {
+  const content = await fs.readFile(path.join(process.cwd(), 'test', file), 'utf8');
+  console.log('content', {
+    content,
+    file
+  });
 
-suite('Python', () => {
-  async function loadConfigurationFile() {
-    const content = await fs.readFile(path.join(process.cwd(), 'test', 'config.py'), 'utf8');
-    console.log('content', {
-      content
-    });
-
-    const config = await native.pythonLoadConfig(
-      content,
-      {
-        file: 'config.py'
-      }
-    );
+  const config = await native.pythonLoadConfig(
+    content,
+    {
+      file
+    }
+  );
 
-    console.log('loaded config', config);
+  console.log(`loaded config ${file}`, config);
 
-    return config;
-  }
+  return config;
+}
 
+suite('Python Config', () => {
   let config: PyConfiguration;
 
   beforeAll(async () => {
-    config = await loadConfigurationFile();
+    config = await loadConfigurationFile('config.py');
   });
 
   test('async checkAuth', async () => {
@@ -91,3 +94,26 @@ suite('Python', () => {
     );
   });
 });
+
+darwinSuite('Scoped Python Config', () => {
+  test('test', async () => {
+    const config = await loadConfigurationFile('scoped-config.py');
+    expect(config).toEqual({
+      schemaPath: 'models',
+      pgSqlPort: 5555,
+      telemetry: false,
+      contextToApiScopes: expect.any(Function),
+      checkAuth: expect.any(Function),
+      queryRewrite: expect.any(Function),
+    });
+
+    if (!config.checkAuth) {
+      throw new Error('checkAuth was not defined in config.py');
+    }
+
+    await config.checkAuth(
+      { requestId: 'test' },
+      'MY_SECRET_TOKEN'
+    );
+  });
+});
diff --git a/packages/cubejs-backend-native/test/scoped-config.py b/packages/cubejs-backend-native/test/scoped-config.py
@@ -0,0 +1,30 @@
+source_code = """
+from cube.conf import settings
+
+settings.schema_path = "models"
+settings.pg_sql_port = 5555
+settings.telemetry = False
+
+def query_rewrite(query, ctx):
+    print('[python] query_rewrite query=', query, ' ctx=', ctx)
+    return query
+
+settings.query_rewrite = query_rewrite
+
+async def check_auth(req, authorization):
+    print('[python] check_auth req=', req, ' authorization=', authorization)
+
+
+settings.check_auth = check_auth
+
+async def context_to_api_scopes():
+    print('[python] context_to_api_scopes')
+    return ['meta', 'data', 'jobs']
+
+settings.context_to_api_scopes = context_to_api_scopes
+"""
+
+__execution_context_globals = {}
+__execution_context_locals = {}
+
+exec(source_code, __execution_context_globals, __execution_context_locals)
diff --git a/packages/cubejs-backend-native/test/templates/scoped-utils.py b/packages/cubejs-backend-native/test/templates/scoped-utils.py
diff --git a/packages/cubejs-backend-native/test/templates/utils.py b/packages/cubejs-backend-native/test/templates/utils.py

Original file line number	Diff line number	Diff line change
`@@ -248,13 +248,13 @@ export interface PyConfiguration {`
`248`	`248`	`contextToApiScopes?: () => Promise<string[]>`
`249`	`249`	`}`
`250`	`250`
`251`		`-export const pythonLoadConfig = async (context: string, options: { file: string }): Promise<PyConfiguration> => {`
	`251`	`+export const pythonLoadConfig = async (content: string, options: { file: string }): Promise<PyConfiguration> => {`
`252`	`252`	`if (isFallbackBuild()) {`
`253`	`253`	`throw new Error('Python is not supported in fallback build');`
`254`	`254`	`}`
`255`	`255`
`256`	`256`	`const native = loadNative();`
`257`		`- const config = await native.pythonLoadConfig(context, options);`
	`257`	`+ const config = await native.pythonLoadConfig(content, options);`
`258`	`258`
`259`	`259`	`if (config.checkAuth) {`
`260`	`260`	`const nativeCheckAuth = config.checkAuth;`