Unified chdb.query with udf_path support

Author: auxten

chdb/__init__.py

@@ -48,21 +48,11 @@ def to_df(r):
 
 
 # wrap _chdb functions
-def query(sql, output_format="CSV"):
+def query(sql, output_format="CSV", path=None, udf_path=None):
     lower_output_format = output_format.lower()
     if lower_output_format == "dataframe":
-        return to_df(_chdb.query(sql, "Arrow"))
+        return to_df(_chdb.query(sql, "Arrow", path=path, udf_path=udf_path))
     elif lower_output_format == 'arrowtable':
-        return to_arrowTable(_chdb.query(sql, "Arrow"))
+        return to_arrowTable(_chdb.query(sql, "Arrow", path=path, udf_path=udf_path))
     else:
-        return _chdb.query(sql, output_format)
-
-
-def query_stateful(sql, output_format="CSV", path=None):
-    lower_output_format = output_format.lower()
-    if lower_output_format == "dataframe":
-        return to_df(_chdb.query_stateful(sql, "Arrow", path))
-    elif lower_output_format == 'arrowtable':
-        return to_arrowTable(_chdb.query_stateful(sql, "Arrow", path))
-    else:
-        return _chdb.query_stateful(sql, output_format, path)
+        return _chdb.query(sql, output_format, path=path, udf_path=udf_path)

programs/local/LocalChdb.cpp

@@ -6,11 +6,12 @@
 extern bool inside_main = true;
 
 
-local_result * queryToBuffer(const std::string & queryStr, const std::string & format = "CSV", const std::string & path = {})
+local_result * queryToBuffer(
+    const std::string & queryStr, const std::string & format = "CSV", const std::string & path = {}, const std::string & udfPath = {})
 {
-    std::vector<std::string> argv = {"clickhouse", "--multiquery"};
+    std::vector<std::string> argv = {"clickhouse", "--", "--multiquery"};
 
-    // if format is "Debug" or "debug", then we will add --verbose and --log-level=trace to argv
+    // If format is "Debug" or "debug", then we will add `--verbose` and `--log-level=trace` to argv
     if (format == "Debug" || format == "debug")
     {
         argv.push_back("--verbose");
@@ -24,7 +25,16 @@ local_result * queryToBuffer(const std::string & queryStr, const std::string & f
         argv.push_back("--output-format=" + format);
     }
 
-    if (!path.empty()) 
+    // If udfPath is not empty, then we will add `--user_scripts_path` and `--user_defined_executable_functions_config` to argv
+    // the path should be a one time thing, so the caller should take care of the temporary files deletion
+    if (!udfPath.empty())
+    {
+        argv.push_back("--user_scripts_path=" + udfPath);
+        argv.push_back("--user_defined_executable_functions_config=" + udfPath + "/*.xml");
+    }
+
+    // If path is not empty, then we will add `--path` to argv. This is used for chdb.Session to support stateful query
+    if (!path.empty())
     {
         // Add path string
         argv.push_back("--path=" + path);
@@ -42,14 +52,10 @@ local_result * queryToBuffer(const std::string & queryStr, const std::string & f
 
 // Pybind11 will take over the ownership of the `query_result` object
 // using smart ptr will cause early free of the object
-query_result * query(const std::string & queryStr, const std::string & format = "CSV")
+query_result *
+query(const std::string & queryStr, const std::string & format = "CSV", const std::string & path = {}, const std::string & udfPath = {})
 {
-    return new query_result(queryToBuffer(queryStr, format));
-}
-
-query_result * query_stateful(const std::string & queryStr, const std::string & format = "CSV", const std::string & path = {})
-{
-    return new query_result(queryToBuffer(queryStr, format, path));
+    return new query_result(queryToBuffer(queryStr, format, path, udfPath));
 }
 
 // The `query_result` and `memoryview_wrapper` will hold `local_result_wrapper` with shared_ptr
@@ -132,9 +138,15 @@ PYBIND11_MODULE(_chdb, m)
         .def("get_memview", &query_result::get_memview);
 
 
-    m.def("query", &query, "Stateless query Clickhouse and return a query_result object");
-
-    m.def("query_stateful", &query_stateful, "Stateful query Clickhouse and return a query_result object");
+    m.def(
+        "query",
+        &query,
+        py::arg("queryStr"),
+        py::arg("format") = "CSV",
+        py::kw_only(),
+        py::arg("path") = "",
+        py::arg("udf_path") = "",
+        "Query chDB and return a query_result object");
 }
 
 #endif // PY_TEST_MAIN