feat(native): Allow to pass variables via TemplateContext from Python…

… to Jinja

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

@@ -364,12 +364,13 @@ export const pythonLoadConfig = async (content: string, options: { fileName: str
 export type PythonCtx = {
     __type: 'PythonCtx'
 } & {
-    [key: string]: Function
+  functions: Record<string, Function>
+  variables: Record<string, Function>
 };
 
 export interface JinjaEngine {
     loadTemplate(templateName: string, templateContent: string): void;
-    renderTemplate(templateName: string, context: unknown, pythonContext: PythonCtx | null): string;
+    renderTemplate(templateName: string, context: unknown, pythonContext: Record<string, any> | null): string;
 }
 
 export class NativeInstance {

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

@@ -164,16 +164,24 @@ class TemplateException(Exception):
 
 class TemplateContext:
     functions: dict[str, Callable]
+    variables: dict[str, Any]
 
     def __init__(self):
         self.functions = {}
+        self.variables = {}
 
     def add_function(self, name, func):
         if not callable(func):
             raise TemplateException("function registration must be used with functions, actual: '%s'" % type(func).__name__)
 
         self.functions[name] = func
 
+    def add_variable(self, name, val):
+        if name in self.functions:
+            raise TemplateException("unable to register variable: name '%s' is already in use for function" % name)
+
+        self.variables[name] = val
+
     def add_filter(self, name, func):
         if not callable(func):
             raise TemplateException("function registration must be used with functions, actual: '%s'" % type(func).__name__)
@@ -194,9 +202,6 @@ def filter(self, func):
         self.add_filter(func.__name__, func)
         return func
 
-    def variable(self, func):
-        raise TemplateException("variable registration is not supported")
-
 class TemplateFunctionRef:
     context: TemplateContext
     attribute: str

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

@@ -1,4 +1,4 @@
-use crate::cross::{CLRepr, CLReprObject, PythonRef};
+use crate::cross::*;
 use crate::python::cube_config::CubeConfigPy;
 use crate::python::python_model::CubePythonModel;
 use crate::python::runtime::py_runtime_init;
@@ -71,10 +71,12 @@ fn python_load_model(mut cx: FunctionContext) -> JsResult<JsPromise> {
 
         let model_module = PyModule::from_code(py, &model_content, &model_file_name, "")?;
         let mut collected_functions = CLReprObject::new();
+        let mut collected_variables = CLReprObject::new();
 
         if model_module.hasattr("template")? {
-            let functions = model_module
-                .getattr("template")?
+            let template = model_module.getattr("template")?;
+
+            let functions = template
                 .getattr("functions")?
                 .downcast::<PyDict>()?;
 
@@ -87,6 +89,18 @@ fn python_load_model(mut cx: FunctionContext) -> JsResult<JsPromise> {
                     );
                 }
             }
+
+            let variables = template
+                .getattr("variables")?
+                .downcast::<PyDict>()?;
+
+            for (local_key, local_value) in variables.iter() {
+                collected_variables.insert(
+                    local_key.to_string(),
+                    CLRepr::from_python_ref(local_value)?,
+                );
+            }
+
             // TODO remove all other ways of defining functions
         } else if model_module.hasattr("__execution_context_locals")? {
             let execution_context_locals = model_module
@@ -128,7 +142,7 @@ fn python_load_model(mut cx: FunctionContext) -> JsResult<JsPromise> {
             }
         };
 
-        Ok(CubePythonModel::new(collected_functions))
+        Ok(CubePythonModel::new(collected_functions, collected_variables))
     });
 
     deferred.settle_with(&channel, move |mut cx| match conf_res {

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

@@ -4,11 +4,12 @@ use crate::cross::{CLRepr, CLReprObject};
 
 pub struct CubePythonModel {
     functions: CLReprObject,
+    variables: CLReprObject,
 }
 
 impl CubePythonModel {
-    pub fn new(functions: CLReprObject) -> Self {
-        Self { functions }
+    pub fn new(functions: CLReprObject, variables: CLReprObject) -> Self {
+        Self { functions, variables }
     }
 }
 
@@ -17,6 +18,10 @@ impl Finalize for CubePythonModel {}
 impl CubePythonModel {
     #[allow(clippy::wrong_self_convention)]
     pub fn to_object<'a, C: Context<'a>>(self, cx: &mut C) -> JsResult<'a, JsValue> {
-        CLRepr::Object(self.functions).into_js(cx)
+        let mut obj = CLReprObject::new();
+        obj.insert("functions".to_string(), CLRepr::Object(self.functions));
+        obj.insert("variables".to_string(), CLRepr::Object(self.variables));
+
+        CLRepr::Object(obj).into_js(cx)
     }
 }

packages/cubejs-backend-native/test/__snapshots__/jinja.test.ts.snap

@@ -250,3 +250,11 @@ exports[`Jinja (new api) render template_error.jinja: template_error.jinja 1`] =
 
 -------------------------------------------------------------------------------]
 `;
+
+exports[`Jinja (new api) render variables.yml.jinja: variables.yml.jinja 1`] = `
+"variables:
+  var1: \\"test string\\"
+  var2: true
+  var3: null
+  var4: {\\"key1\\":\\"value1\\",\\"key2\\":true}"
+`;

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

@@ -38,7 +38,11 @@ function testTemplateBySnapshot(engine: JinjaEngine, templateName: string, ctx:
 
 function testTemplateWithPythonCtxBySnapshot(engine: JinjaEngine, templateName: string, ctx: unknown, utilsFile: string) {
   test(`render ${templateName}`, async () => {
-    const actual = engine.renderTemplate(templateName, ctx, await loadPythonCtxFromUtils(utilsFile));
+    const pyCtx = await loadPythonCtxFromUtils(utilsFile);
+    const actual = engine.renderTemplate(templateName, ctx, {
+      ...pyCtx.variables,
+      ...pyCtx.functions,
+    });
 
     expect(actual).toMatchSnapshot(templateName);
   });
@@ -60,7 +64,7 @@ suite('Python model', () => {
   it('load jinja-instance.py', async () => {
     const pythonModule = await loadPythonCtxFromUtils('jinja-instance.py');
 
-    expect(pythonModule).toEqual({
+    expect(pythonModule.functions).toEqual({
       load_data: expect.any(Object),
       load_data_sync: expect.any(Object),
       arg_bool: expect.any(Object),
@@ -75,14 +79,21 @@ suite('Python model', () => {
       new_safe_string: expect.any(Object),
       load_class_model: expect.any(Object),
     });
+
+    expect(pythonModule.variables).toEqual({
+      var1: 'test string',
+      var2: true,
+      var3: undefined,
+      var4: expect.any(Object),
+    });
   });
 });
 
 darwinSuite('Scope Python model', () => {
   it('load scoped-utils.py', async () => {
     const pythonModule = await loadPythonCtxFromUtils('scoped-utils.py');
 
-    expect(pythonModule).toEqual({
+    expect(pythonModule.functions).toEqual({
       load_data: expect.any(Object),
       load_data_sync: expect.any(Object),
       arg_bool: expect.any(Object),
@@ -113,6 +124,7 @@ function createTestSuite(utilsFile: string) {
       loadTemplateFile(jinjaEngine, 'data-model.yml.jinja');
       loadTemplateFile(jinjaEngine, 'arguments-test.yml.jinja');
       loadTemplateFile(jinjaEngine, 'python.yml');
+      loadTemplateFile(jinjaEngine, 'variables.yml.jinja');
 
       for (let i = 1; i < 9; i++) {
         loadTemplateFile(jinjaEngine, `0${i}.yml.jinja`);
@@ -139,6 +151,7 @@ function createTestSuite(utilsFile: string) {
     testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'data-model.yml.jinja', {}, utilsFile);
     testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'arguments-test.yml.jinja', {}, utilsFile);
     testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'python.yml', {}, utilsFile);
+    testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'variables.yml.jinja', {}, utilsFile);
 
     testLoadBrokenTemplateBySnapshot(jinjaEngine, 'template_error.jinja');
 

packages/cubejs-backend-native/test/templates/jinja-instance.py

@@ -1,6 +1,10 @@
 from cube import (TemplateContext, SafeString)
 
 template = TemplateContext()
+template.add_variable('var1', "test string")
+template.add_variable('var2', True)
+template.add_variable('var3', None)
+template.add_variable('var4', {'key1': 'value1', 'key2': True})
 
 @template.function
 def arg_sum_integers(a, b):

packages/cubejs-backend-native/test/templates/variables.yml.jinja

@@ -0,0 +1,5 @@
+variables:
+  var1: {{ var1 }}
+  var2: {{ var2 }}
+  var3: {{ var3 }}
+  var4: {{ var4 }}

packages/cubejs-schema-compiler/src/compiler/DataSchemaCompiler.js

@@ -68,7 +68,10 @@ export class DataSchemaCompiler {
 
       return {
         fileName: file.fileName,
-        exports
+        exports: {
+          ...exports.variables,
+          ...exports.functions,
+        }
       };
     }));