From a69b5310aaba93dc8482a56d155886710b506308 Mon Sep 17 00:00:00 2001
From: sundyli <543950155@qq.com>
Date: Fri, 6 Dec 2024 11:05:53 +0800
Subject: [PATCH] fix(cli): stats are missing for update queries (#527)

---
 cli/src/session.rs              | 16 --------------
 cli/tests/http/06-update.result |  1 +
 cli/tests/http/06-update.sh     | 37 +++++++++++++++++++++++++++++++++
 core/src/response.rs            | 15 +++++++++++++
 driver/src/rest_api.rs          |  9 ++++++--
 5 files changed, 60 insertions(+), 18 deletions(-)
 create mode 100644 cli/tests/http/06-update.result
 create mode 100644 cli/tests/http/06-update.sh

diff --git a/cli/src/session.rs b/cli/src/session.rs
index 1701501b3..ca4a52132 100644
--- a/cli/src/session.rs
+++ b/cli/src/session.rs
@@ -473,22 +473,6 @@ impl Session {
                 let _ = self.conn.exec(query).await?;
                 Ok(None)
             }
-            QueryKind::Update => {
-                let affected = self.conn.exec(query).await?;
-                if is_repl {
-                    if affected > 0 {
-                        eprintln!(
-                            "{} rows affected in ({:.3} sec)",
-                            affected,
-                            start.elapsed().as_secs_f64()
-                        );
-                    } else {
-                        eprintln!("processed in ({:.3} sec)", start.elapsed().as_secs_f64());
-                    }
-                    eprintln!();
-                }
-                Ok(Some(ServerStats::default()))
-            }
             other => {
                 let replace_newline = !if self.settings.replace_newline {
                     false
diff --git a/cli/tests/http/06-update.result b/cli/tests/http/06-update.result
new file mode 100644
index 000000000..467c029b3
--- /dev/null
+++ b/cli/tests/http/06-update.result
@@ -0,0 +1 @@
+305104
diff --git a/cli/tests/http/06-update.sh b/cli/tests/http/06-update.sh
new file mode 100644
index 000000000..a5cd37d63
--- /dev/null
+++ b/cli/tests/http/06-update.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+
+cat <<SQL | ${BENDSQL}
+CREATE TABLE t3 (
+    category_id INT NOT NULL,
+    region_id INT NOT NULL,
+    user_id BIGINT NOT NULL,
+    product_id INT NOT NULL,
+    transaction_date DATETIME NOT NULL,
+    transaction_amount DECIMAL(10,2) NOT NULL,
+    discount_amount DECIMAL(10,2) NOT NULL,
+    is_returned BOOLEAN NOT NULL DEFAULT 0,
+    inventory_count INT NOT NULL,
+    last_updated TIMESTAMP NOT NULL
+);
+
+create table r like t3 engine = random;
+
+insert into t3 select * from r limit 305104;
+SQL
+
+
+a=`${BENDSQL} --stats  --query="""
+update t3 set user_id = user_id + 1 where inventory_count % 10 >= 0;
+""" 2>&1 | grep -oE '([0-9]+) rows written' | grep -oE '([0-9]+)'`
+
+
+b=`${BENDSQL} --stats  --query="""
+update t3 set user_id = user_id + 1 where inventory_count % 10 < 0;
+""" 2>&1 | grep -oE '([0-9]+) rows written' | grep -oE '([0-9]+)'`
+
+echo "$[a+b]"
+
+
+cat <<SQL | ${BENDSQL}
+DROP TABLE IF EXISTS t3;
+SQL
diff --git a/core/src/response.rs b/core/src/response.rs
index b789c925c..048964f58 100644
--- a/core/src/response.rs
+++ b/core/src/response.rs
@@ -32,6 +32,21 @@ pub struct Progresses {
     pub total_scan: Option<ProgressValues>,
 }
 
+impl Progresses {
+    pub fn has_progress(&self) -> bool {
+        self.scan_progress.bytes > 0
+            || self.scan_progress.rows > 0
+            || self.write_progress.bytes > 0
+            || self.write_progress.rows > 0
+            || self.result_progress.bytes > 0
+            || self.result_progress.rows > 0
+            || self
+                .total_scan
+                .as_ref()
+                .map_or(false, |v| v.bytes > 0 || v.rows > 0)
+    }
+}
+
 #[derive(Debug, Deserialize)]
 pub struct ProgressValues {
     pub rows: usize,
diff --git a/driver/src/rest_api.rs b/driver/src/rest_api.rs
index 9a8794cfa..471c0f280 100644
--- a/driver/src/rest_api.rs
+++ b/driver/src/rest_api.rs
@@ -204,7 +204,8 @@ impl<'o> RestAPIConnection {
     }
 
     async fn wait_for_schema(&self, resp: QueryResponse) -> Result<QueryResponse> {
-        if !resp.data.is_empty() || !resp.schema.is_empty() {
+        if !resp.data.is_empty() || !resp.schema.is_empty() || resp.stats.progresses.has_progress()
+        {
             return Ok(resp);
         }
         let node_id = resp.node_id.clone();
@@ -218,7 +219,11 @@ impl<'o> RestAPIConnection {
                 .client
                 .query_page(&result.id, &next_uri, &node_id)
                 .await?;
-            if !result.data.is_empty() || !result.schema.is_empty() {
+
+            if !result.data.is_empty()
+                || !result.schema.is_empty()
+                || result.stats.progresses.has_progress()
+            {
                 break;
             }
         }