z

Author: hantmac

driver/src/rest_api.rs

@@ -64,9 +64,7 @@ impl IConnection for RestAPIConnection {
 
     async fn exec(&self, sql: &str) -> Result<i64> {
         info!("exec: {}", sql);
-        let page = self.client.query_all(sql).await?;
-        let affected_rows = parse_affected_rows_from_page(&page)?;
-        Ok(affected_rows)
+        self.calculate_affected_rows_from_iter(sql).await
     }
 
     async fn kill_query(&self, query_id: &str) -> Result<()> {
@@ -178,13 +176,13 @@ impl IConnection for RestAPIConnection {
     }
 }
 
-impl<'o> RestAPIConnection {
+impl RestAPIConnection {
     pub async fn try_create(dsn: &str, name: String) -> Result<Self> {
         let client = APIClient::new(dsn, Some(name)).await?;
         Ok(Self { client })
     }
 
-    fn default_file_format_options() -> BTreeMap<&'o str, &'o str> {
+    fn default_file_format_options() -> BTreeMap<&'static str, &'static str> {
         vec![
             ("type", "CSV"),
             ("field_delimiter", ","),
@@ -195,9 +193,68 @@ impl<'o> RestAPIConnection {
         .collect()
     }
 
-    fn default_copy_options() -> BTreeMap<&'o str, &'o str> {
+    fn default_copy_options() -> BTreeMap<&'static str, &'static 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;
+
+        use tokio_stream::StreamExt;
+        // 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;
+                // Continue counting the rest
+                while let Some(row_result) = rows.next().await {
+                    row_result?;
+                    count += 1;
+                }
+            }
+        }
+
+        Ok(count)
+    }
 }
 
 pub struct RestAPIRows<T> {
@@ -288,49 +345,3 @@ impl FromRowStats for RawRowWithStats {
         Ok(RawRowWithStats::Row(RawRow::new(rows, row)))
     }
 }
-
-fn parse_affected_rows_from_page(page: &databend_client::Page) -> Result<i64> {
-    if page.schema.is_empty() {
-        return Ok(0);
-    }
-
-    let first_field = &page.schema[0];
-    if !first_field.name.contains("number of rows") {
-        return Ok(0);
-    }
-
-    if page.data.is_empty() || page.data[0].is_empty() {
-        return Ok(0);
-    }
-
-    match &page.data[0][0] {
-        Some(value_str) => parse_row_count_string(value_str).map_err(Error::InvalidResponse),
-        None => Ok(0),
-    }
-}
-
-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
-    ))
-}