z

databendlabs · Dec 26, 2024 · bd18948 · bd18948
1 parent 49a50a4
commit bd18948
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 20 deletions.
diff --git a/bindings/python/src/blocking.rs b/bindings/python/src/blocking.rs
@@ -197,28 +197,26 @@ impl BlockingDatabendCursor {
         &'p mut self,
         py: Python<'p>,
         operation: String,
-        parameters: Option<Vec<Bound<'p, PyAny>>>,
+        parameters: Option<Bound<'p, PyAny>>,
     ) -> PyResult<PyObject> {
         self.reset();
         let conn = self.conn.clone();
-        if let Some(parameters) = parameters {
-            if let Some(first) = parameters.first() {
-                if first.downcast::<PyList>().is_ok() || first.downcast::<PyTuple>().is_ok() {
-                    let bytes = format_csv(parameters)?;
-                    let size = bytes.len() as u64;
-                    let reader = Box::new(std::io::Cursor::new(bytes));
-                    let stats = wait_for_future(py, async move {
-                        conn.load_data(&operation, reader, size, None, None)
-                            .await
-                            .map_err(DriverError::new)
-                    })?;
-                    let result = stats.write_rows.into_pyobject(py)?;
-                    return Ok(result.into());
-                } else {
-                    return Err(PyAttributeError::new_err(
-                        "Invalid parameter type, expected list or tuple",
-                    ));
-                }
+        if let Some(param) = parameters {
+            if param.downcast::<PyList>().is_ok() || param.downcast::<PyTuple>().is_ok() {
+                let bytes = format_csv([param].to_vec())?;
+                let size = bytes.len() as u64;
+                let reader = Box::new(std::io::Cursor::new(bytes));
+                let stats = wait_for_future(py, async move {
+                    conn.load_data(&operation, reader, size, None, None)
+                        .await
+                        .map_err(DriverError::new)
+                })?;
+                let result = stats.write_rows.into_pyobject(py)?;
+                return Ok(result.into());
+            } else {
+                return Err(PyAttributeError::new_err(
+                    "Invalid parameter type, expected list or tuple",
+                ));
             }
         }
         // fetch first row after execute
@@ -236,6 +234,35 @@ impl BlockingDatabendCursor {
         Ok(py.None())
     }
 
+    pub fn executemany<'p>(
+        &'p mut self,
+        py: Python<'p>,
+        operation: String,
+        parameters: Vec<Bound<'p, PyAny>>,
+    ) -> PyResult<PyObject> {
+        self.reset();
+        let conn = self.conn.clone();
+        if let Some(param) = parameters.first() {
+            if param.downcast::<PyList>().is_ok() || param.downcast::<PyTuple>().is_ok() {
+                let bytes = format_csv(parameters)?;
+                let size = bytes.len() as u64;
+                let reader = Box::new(std::io::Cursor::new(bytes));
+                let stats = wait_for_future(py, async move {
+                    conn.load_data(&operation, reader, size, None, None)
+                        .await
+                        .map_err(DriverError::new)
+                })?;
+                let result = stats.write_rows.into_pyobject(py)?;
+                return Ok(result.into());
+            } else {
+                return Err(PyAttributeError::new_err(
+                    "Invalid parameter type, expected list or tuple",
+                ));
+            }
+        }
+        Ok(py.None())
+    }
+
     pub fn fetchone(&mut self, py: Python) -> PyResult<Option<Row>> {
         if let Some(row) = self.buffer.pop() {
             return Ok(Some(row));

diff --git a/bindings/python/tests/blocking/steps/binding.py b/bindings/python/tests/blocking/steps/binding.py
@@ -154,7 +154,7 @@ def _(context):
     ret = context.cursor.fetchone()
     assert ret is None, f"ret: {ret}"
 
-    context.cursor.execute("INSERT INTO test VALUES", values)
+    context.cursor.executemany("INSERT INTO test VALUES", values)
     ret = context.cursor.fetchall()
     assert ret == expected, f"ret: {ret}"