feat(native): Improve safety for Python into CLRepr conversion (#9984)

ovr · web-flow · commit 644103a8d0c5 · 2025-09-22T17:41:48.000+02:00
diff --git a/packages/cubejs-backend-native/src/cross/clrepr_python.rs b/packages/cubejs-backend-native/src/cross/clrepr_python.rs
@@ -1,9 +1,10 @@
 use crate::cross::clrepr::CLReprObject;
 use crate::cross::{CLRepr, CLReprObjectKind, StringType};
-use pyo3::exceptions::{PyNotImplementedError, PyTypeError};
+use crate::python::utils::PyAnyHelpers;
+use pyo3::exceptions::{PyException, PyNotImplementedError, PyTypeError};
 use pyo3::types::{
-    PyBool, PyComplex, PyDate, PyDict, PyFloat, PyFrame, PyFunction, PyInt, PyList, PySequence,
-    PySet, PyString, PyTraceback, PyTuple,
+    PyBool, PyComplex, PyDate, PyDateTime, PyDelta, PyDict, PyFloat, PyFrame, PyFrozenSet,
+    PyFunction, PyInt, PyList, PySequence, PySet, PyString, PyTraceback, PyTuple,
 };
 use pyo3::{Py, PyAny, PyErr, PyObject, Python, ToPyObject};
 
@@ -12,7 +13,7 @@ pub enum PythonRef {
     PyObject(PyObject),
     PyFunction(Py<PyFunction>),
     /// Special type to transfer functions through JavaScript
-    /// In JS it's an external object. It's not the same as Function.
+    /// In JS it's an external object.
     PyExternalFunction(Py<PyFunction>),
 }
 
@@ -87,10 +88,28 @@ impl CLRepr {
             return Err(PyErr::new::<PyTypeError, _>(
                 "Unable to represent PyComplex type as CLR from Python".to_string(),
             ));
+        } else if v.get_type().is_subclass_of::<PyDateTime>()? {
+            return Err(PyErr::new::<PyTypeError, _>(
+                "Unable to represent PyDateTime type as CLR from Python".to_string(),
+            ));
         } else if v.get_type().is_subclass_of::<PyDate>()? {
             return Err(PyErr::new::<PyTypeError, _>(
                 "Unable to represent PyDate type as CLR from Python".to_string(),
             ));
+        } else if v.get_type().is_subclass_of::<PyFrozenSet>()? {
+            let set = v.downcast::<PyFrozenSet>()?;
+
+            return Err(PyErr::new::<PyTypeError, _>(format!(
+                "Unable to represent PyFrozenSet type as CLR from Python, value: {:?}",
+                set
+            )));
+        } else if v.get_type().is_subclass_of::<PyException>()? {
+            let exception = v.downcast::<PyException>()?;
+
+            return Err(PyErr::new::<PyTypeError, _>(format!(
+                "Unable to represent PyException type as CLR from Python, value: {:?}",
+                exception
+            )));
         } else if v.get_type().is_subclass_of::<PyFrame>()? {
             let frame = v.downcast::<PyFrame>()?;
 
@@ -105,17 +124,22 @@ impl CLRepr {
                 "Unable to represent PyTraceback type as CLR from Python, value: {:?}",
                 trb
             )));
-        } else {
-            let is_sequence = unsafe { pyo3::ffi::PySequence_Check(v.as_ptr()) == 1 };
-            if is_sequence {
-                let seq = v.downcast::<PySequence>()?;
-
-                return Err(PyErr::new::<PyTypeError, _>(format!(
-                    "Unable to represent PySequence type as CLR from Python, value: {:?}",
-                    seq
-                )));
-            }
+        } else if v.get_type().is_subclass_of::<PyDelta>()? {
+            let delta = v.downcast::<PyDelta>()?;
+
+            return Err(PyErr::new::<PyTypeError, _>(format!(
+                "Unable to represent PyDelta type as CLR from Python, value: {:?}",
+                delta
+            )));
+        } else if v.is_sequence()? {
+            let seq = v.downcast::<PySequence>()?;
 
+            return Err(PyErr::new::<PyTypeError, _>(format!(
+                "Unable to represent PySequence type as CLR from Python, value: {:?}",
+                seq
+            )));
+        } else {
+            // Fallback to PyObject, it will lead to throw error in the JS side
             Self::PythonRef(PythonRef::PyObject(v.into()))
         })
     }
diff --git a/packages/cubejs-backend-native/src/python/runtime.rs b/packages/cubejs-backend-native/src/python/runtime.rs
@@ -1,5 +1,6 @@
 use crate::cross::CLRepr;
 use crate::python::neon_py::*;
+use crate::python::utils::PyAnyHelpers;
 use crate::tokio_runtime_node;
 use cubesql::CubeError;
 use log::{error, trace};
@@ -143,8 +144,7 @@ impl PyRuntime {
             let py_args = PyTuple::new(py, prep_tuple);
             let call_res = fun.call(py, py_args, py_kwargs)?;
 
-            let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
-            if is_coroutine {
+            if call_res.is_coroutine()? {
                 let fut = pyo3_asyncio::tokio::into_future(call_res.as_ref(py))?;
                 Ok(PyScheduledFunResult::Poll(Box::pin(fut)))
             } else {
diff --git a/packages/cubejs-backend-native/src/python/utils.rs b/packages/cubejs-backend-native/src/python/utils.rs
@@ -1,7 +1,9 @@
 use crate::cross::*;
-use pyo3::exceptions::PyNotImplementedError;
+use pyo3::exceptions::{PyNotImplementedError, PySystemError};
 use pyo3::prelude::*;
 use pyo3::types::{PyFunction, PyString, PyTuple};
+use pyo3::{ffi, AsPyPointer};
+use std::ffi::c_int;
 
 pub fn python_fn_call_sync(py_fun: &Py<PyFunction>, arguments: Vec<CLRepr>) -> PyResult<CLRepr> {
     Python::with_gil(|py| {
@@ -15,8 +17,7 @@ pub fn python_fn_call_sync(py_fun: &Py<PyFunction>, arguments: Vec<CLRepr>) -> P
 
         let call_res = py_fun.call1(py, tuple)?;
 
-        let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
-        if is_coroutine {
+        if call_res.is_coroutine()? {
             Err(PyErr::new::<PyNotImplementedError, _>(
                 "Calling function with async response is not supported",
             ))
@@ -38,8 +39,7 @@ pub fn python_obj_call_sync(py_fun: &PyObject, arguments: Vec<CLRepr>) -> PyResu
 
         let call_res = py_fun.call1(py, tuple)?;
 
-        let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
-        if is_coroutine {
+        if call_res.is_coroutine()? {
             Err(PyErr::new::<PyNotImplementedError, _>(
                 "Calling object with async response is not supported",
             ))
@@ -68,8 +68,7 @@ where
 
         let call_res = py_fun.call_method1(py, name, tuple)?;
 
-        let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
-        if is_coroutine {
+        if call_res.is_coroutine()? {
             Err(PyErr::new::<PyNotImplementedError, _>(
                 "Calling object method with async response is not supported",
             ))
@@ -78,3 +77,52 @@ where
         }
     })
 }
+
+pub trait PyAnyHelpers {
+    fn is_sequence(&self) -> PyResult<bool>;
+
+    fn is_coroutine(&self) -> PyResult<bool>;
+}
+
+pub(crate) fn internal_error_on_minusone(result: c_int) -> PyResult<()> {
+    if result != -1 {
+        Ok(())
+    } else {
+        Err(PyErr::new::<PySystemError, _>(
+            "Error on call via ffi, result is -1",
+        ))
+    }
+}
+
+impl<T> PyAnyHelpers for T
+where
+    T: AsPyPointer,
+{
+    fn is_sequence(&self) -> PyResult<bool> {
+        let ptr = self.as_ptr();
+        if ptr.is_null() {
+            return Err(PyErr::new::<PySystemError, _>(
+                "Unable to verify that object is sequence, because ptr is null",
+            ));
+        }
+
+        let v = unsafe { ffi::PySequence_Check(ptr) };
+        internal_error_on_minusone(v)?;
+
+        Ok(v != 0)
+    }
+
+    fn is_coroutine(&self) -> PyResult<bool> {
+        let ptr = self.as_ptr();
+        if ptr.is_null() {
+            return Err(PyErr::new::<PySystemError, _>(
+                "Unable to verify that object is coroutine, because ptr is null",
+            ));
+        }
+
+        let v = unsafe { ffi::PyCoro_CheckExact(ptr) };
+        internal_error_on_minusone(v)?;
+
+        Ok(v != 0)
+    }
+}
