Revert "impl affect_rows in rest_api.rs and use query_iter"

This reverts commit e629133.

Author: hantmac

bindings/python/src/blocking.rs

@@ -294,12 +294,11 @@ impl BlockingDatabendCursor {
                     .await
                     .map_err(DriverError::new)
             })?;
-
             self.rowcount = affected_rows;
             return Ok(py.None());
         }
 
-        // For SELECT, use query_iter as before
+        //  for select, use query_iter
         let (first, rows) = wait_for_future(py, async move {
             let mut rows = conn.query_iter(&operation, params).await?;
             let first = rows.next().await.transpose()?;

core/src/pages.rs

@@ -47,6 +47,58 @@ impl Page {
         }
         self.stats = p.stats;
     }
+
+    pub fn affected_rows(&self) -> Result<i64, String> {
+        if self.schema.is_empty() {
+            return Ok(0);
+        }
+
+        let first_field = &self.schema[0];
+
+        if !first_field.name.contains("number of rows") {
+            return Ok(0);
+        }
+
+        if self.data.is_empty() || self.data[0].is_empty() {
+            return Ok(0);
+        }
+
+        match &self.data[0][0] {
+            Some(value_str) => self.parse_row_count_string(value_str),
+            None => Ok(0),
+        }
+    }
+
+    fn parse_row_count_string(&self, value_str: &str) -> Result<i64, String> {
+        let trimmed = value_str.trim();
+
+        if trimmed.is_empty() {
+            return Ok(0);
+        }
+
+        if let Ok(count) = trimmed.parse::<i64>() {
+            return Ok(count);
+        }
+
+        if let Ok(count) = serde_json::from_str::<i64>(trimmed) {
+            return Ok(count);
+        }
+
+        let unquoted = trimmed.trim_matches('"');
+        if let Ok(count) = unquoted.parse::<i64>() {
+            return Ok(count);
+        }
+
+        Err(format!(
+            "failed to parse affected rows from: '{}'",
+            value_str
+        ))
+    }
+
+    ///the schema can be `number of rows inserted`, `number of rows deleted`, `number of rows updated` when sql start with  `insert`, `delete`, `update`
+    pub fn has_affected_rows(&self) -> bool {
+        !self.schema.is_empty() && self.schema[0].name.contains("number of rows")
+    }
 }
 
 type PageFut = Pin<Box<dyn Future<Output = Result<QueryResponse>> + Send>>;

driver/src/rest_api.rs

@@ -24,7 +24,6 @@ use log::info;
 use tokio::fs::File;
 use tokio::io::BufReader;
 use tokio_stream::Stream;
-use tokio_stream::StreamExt;
 
 use databend_client::APIClient;
 use databend_client::Pages;
@@ -65,8 +64,11 @@ impl IConnection for RestAPIConnection {
 
     async fn exec(&self, sql: &str) -> Result<i64> {
         info!("exec: {}", sql);
-        // Use the new affected_rows method that internally uses query_iter
-        self.calculate_affected_rows_from_iter(sql).await
+        let page = self.client.query_all(sql).await?;
+
+        let affected_rows = page.affected_rows().map_err(Error::InvalidResponse)?;
+
+        Ok(affected_rows)
     }
 
     async fn kill_query(&self, query_id: &str) -> Result<()> {
@@ -198,60 +200,6 @@ impl<'o> RestAPIConnection {
     fn default_copy_options() -> BTreeMap<&'o str, &'o str> {
         vec![("purge", "true")].into_iter().collect()
     }
-
-    fn parse_row_count_string(value_str: &str) -> Result<i64, String> {
-        let trimmed = value_str.trim();
-
-        if trimmed.is_empty() {
-            return Ok(0);
-        }
-
-        if let Ok(count) = trimmed.parse::<i64>() {
-            return Ok(count);
-        }
-
-        if let Ok(count) = serde_json::from_str::<i64>(trimmed) {
-            return Ok(count);
-        }
-
-        let unquoted = trimmed.trim_matches('"');
-        if let Ok(count) = unquoted.parse::<i64>() {
-            return Ok(count);
-        }
-
-        Err(format!(
-            "failed to parse affected rows from: '{}'",
-            value_str
-        ))
-    }
-
-    async fn calculate_affected_rows_from_iter(&self, sql: &str) -> Result<i64> {
-        let mut rows = IConnection::query_iter(self, sql).await?;
-        let mut count = 0i64;
-
-        // Get the first row to check if it has affected rows info
-        if let Some(first_row) = rows.next().await {
-            let row = first_row?;
-            let schema = row.schema();
-
-            // Check if this is an affected rows response
-            if !schema.fields().is_empty() && schema.fields()[0].name.contains("number of rows") {
-                let values = row.values();
-                if !values.is_empty() {
-                    let value = &values[0];
-                    let s: String = value.clone().try_into().map_err(|e| {
-                        Error::InvalidResponse(format!("Failed to convert value to string: {}", e))
-                    })?;
-                    count = Self::parse_row_count_string(&s).map_err(Error::InvalidResponse)?;
-                }
-            } else {
-                // If it's not affected rows info, count normally
-                count = -1;
-            }
-        }
-
-        Ok(count)
-    }
 }
 
 pub struct RestAPIRows<T> {