feat(schema): Support passing arguments for @context_func in jinja templates (#6882)

Author: ovr

packages/cubejs-backend-native/Cargo.lock

@@ -28,7 +28,7 @@ libc = "0.2"
 findshlibs = "0.10.2"
 convert_case = "0.6.0"
 # python
-minijinja = { git = "https://github.com/mitsuhiko/minijinja.git", rev = "efdb9bee8b893b423157035631cfe3ae694c45cd", features = ["json", "loader"], optional = true }
+minijinja = { version = "1", features = ["json", "loader"], optional = true }
 pyo3 = { version = "0.18.3", features = [], optional = true }
 pyo3-asyncio = { version = "0.18.0", features = ["tokio-runtime", "attributes"], optional = true }
 

packages/cubejs-backend-native/Cargo.toml

@@ -7,7 +7,7 @@ use pyo3::exceptions::{PyNotImplementedError, PyTypeError};
 use pyo3::types::{PyBool, PyDict, PyFloat, PyFunction, PyInt, PyList, PyString};
 use pyo3::{Py, PyAny, PyErr, PyObject, Python, ToPyObject};
 use std::cell::RefCell;
-use std::collections::hash_map::{IntoIter, Iter};
+use std::collections::hash_map::{IntoIter, Iter, Keys};
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -34,6 +34,10 @@ impl CLReprObject {
     pub fn iter(&self) -> Iter<String, CLRepr> {
         self.0.iter()
     }
+
+    pub fn keys(&self) -> Keys<'_, String, CLRepr> {
+        self.0.keys()
+    }
 }
 
 impl std::fmt::Debug for CLReprObject {

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

@@ -2,9 +2,10 @@ use crate::python::cross::{CLRepr, CLReprKind, CLReprObject};
 use crate::python::runtime::py_runtime;
 use crate::tokio_runtime;
 use minijinja as mj;
-use minijinja::value::{Object, ObjectKind, SeqObject, StructObject, Value};
+use minijinja::value::{Object, ObjectKind, SeqObject, StructObject, Value, ValueKind};
 use pyo3::types::PyFunction;
 use pyo3::Py;
+use std::convert::TryInto;
 use std::sync::Arc;
 
 struct JinjaPythonFunction {
@@ -41,16 +42,15 @@ impl Object for JinjaPythonFunction {
     }
 
     fn call(&self, _state: &mj::State, args: &[Value]) -> Result<Value, mj::Error> {
-        if args.len() > 0 {
-            return Err(mj::Error::new(
-                mj::ErrorKind::EvalBlock,
-                "Passing argument to python functions is not supported".to_string(),
-            ));
+        let mut arguments = Vec::with_capacity(args.len());
+
+        for arg in args {
+            arguments.push(from_minijinja_value(arg)?);
         }
 
         let py_runtime = py_runtime()
             .map_err(|err| mj::Error::new(mj::ErrorKind::EvalBlock, format!("Error: {}", err)))?;
-        let call_future = py_runtime.call_async(self.inner.clone(), vec![]);
+        let call_future = py_runtime.call_async(self.inner.clone(), arguments);
 
         let tokio = tokio_runtime()
             .map_err(|err| mj::Error::new(mj::ErrorKind::EvalBlock, format!("Error: {}", err)))?;
@@ -91,11 +91,8 @@ impl StructObject for JinjaDynamicObject {
         self.inner.get(name).map(|v| to_minijinja_value(v.clone()))
     }
 
-    fn fields(&self) -> Vec<Arc<String>> {
-        self.inner
-            .iter()
-            .map(|(k, _)| Arc::new(k.clone()))
-            .collect()
+    fn fields(&self) -> Vec<Arc<str>> {
+        self.inner.keys().map(|x| x.to_string().into()).collect()
     }
 
     fn field_count(&self) -> usize {
@@ -180,6 +177,37 @@ impl Object for JinjaSequenceObject {
     }
 }
 
+pub fn from_minijinja_value(from: &mj::value::Value) -> Result<CLRepr, mj::Error> {
+    match from.kind() {
+        ValueKind::Undefined | ValueKind::None => Ok(CLRepr::Null),
+        ValueKind::Bool => Ok(CLRepr::Bool(from.is_true())),
+        ValueKind::Number => {
+            if let Ok(rv) = TryInto::<i64>::try_into(from.clone()) {
+                Ok(CLRepr::Int(rv))
+            } else if let Ok(rv) = TryInto::<f64>::try_into(from.clone()) {
+                Ok(CLRepr::Float(rv))
+            } else {
+                Err(mj::Error::new(
+                    mj::ErrorKind::InvalidOperation,
+                    format!("Converting from {:?} to python is not supported", from),
+                ))
+            }
+        }
+        ValueKind::String => {
+            if from.is_safe() {
+                // TODO: Danger?
+                Ok(CLRepr::String(from.as_str().unwrap().to_string()))
+            } else {
+                Ok(CLRepr::String(from.as_str().unwrap().to_string()))
+            }
+        }
+        other => Err(mj::Error::new(
+            mj::ErrorKind::InvalidOperation,
+            format!("Converting from {:?} to python is not supported", other),
+        )),
+    }
+}
+
 pub fn to_minijinja_value(from: CLRepr) -> Value {
     match from {
         CLRepr::Array(inner) => Value::from_seq_object(JinjaSequenceObject { inner }),

packages/cubejs-backend-native/src/python/template/mj_value.rs

@@ -174,6 +174,16 @@ exports[`Jinja render 08.yml.jinja: 08.yml.jinja 1`] = `
     data_source: postgres }"
 `;
 
+exports[`Jinja render arguments-test.yml.jinja: arguments-test.yml.jinja 1`] = `
+"test:
+  arg_sum_integers_int_int: 2
+  arg_sum_integers_int_float: 4.140000000000001
+  arg_bool_true: 1
+  arg_bool_false: 0
+  arg_str: hello world
+  arg_null: none"
+`;
+
 exports[`Jinja render data-model.yml.jinja: data-model.yml.jinja 1`] = `
 "cubes:
 

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

@@ -56,6 +56,10 @@ suite('Python models', () => {
     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),
     });
   });
 });
@@ -69,6 +73,7 @@ suite('Jinja', () => {
     loadTemplateFile(jinjaEngine, '.utils.jinja');
     loadTemplateFile(jinjaEngine, 'dump_context.yml.jinja');
     loadTemplateFile(jinjaEngine, 'data-model.yml.jinja');
+    loadTemplateFile(jinjaEngine, 'arguments-test.yml.jinja');
 
     for (let i = 1; i < 9; i++) {
       loadTemplateFile(jinjaEngine, `0${i}.yml.jinja`);
@@ -90,6 +95,7 @@ suite('Jinja', () => {
     }
   });
   testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'data-model.yml.jinja', {});
+  testTemplateWithPythonCtxBySnapshot(jinjaEngine, 'arguments-test.yml.jinja', {});
 
   testLoadBrokenTemplateBySnapshot(jinjaEngine, 'template_error.jinja');
 

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

@@ -0,0 +1,7 @@
+test:
+  arg_sum_integers_int_int: {{ arg_sum_integers(1, 1) }}
+  arg_sum_integers_int_float: {{ arg_sum_integers(1, 3.14) }}
+  arg_bool_true: {{ arg_bool(true) }}
+  arg_bool_false: {{ arg_bool(false) }}
+  arg_str: {{ arg_str("hello world") }}
+  arg_null: {{ arg_null(null) }}

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

@@ -1,5 +1,21 @@
 from cube import context_func
 
+@context_func
+def arg_sum_integers(a, b):
+  return a + b
+
+@context_func
+def arg_bool(a):
+  return a + 0
+
+@context_func
+def arg_str(a):
+  return a
+
+@context_func
+def arg_null(a):
+  return a
+
 @context_func
 def load_data_sync():
    client = MyApiClient("google.com")