From 841387ce04b94e5debf0b61268c0fd5c217a5ffb Mon Sep 17 00:00:00 2001
From: sundyli <543950155@qq.com>
Date: Thu, 27 Jul 2023 18:53:38 -0700
Subject: [PATCH] feat(cli): send write progress (#174)

---
 cli/src/display.rs | 17 ++++++++++++++---
 cli/src/session.rs |  6 +++---
 licenserc.toml     |  3 +++
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/cli/src/display.rs b/cli/src/display.rs
index cad2bee4b..d69ef80a9 100644
--- a/cli/src/display.rs
+++ b/cli/src/display.rs
@@ -31,6 +31,7 @@ use crate::{
     ast::format_query,
     config::{OutputFormat, Settings},
     helper::CliHelper,
+    session::QueryKind,
 };
 
 #[async_trait::async_trait]
@@ -42,6 +43,7 @@ pub trait ChunkDisplay {
 pub struct FormatDisplay<'a> {
     settings: &'a Settings,
     query: &'a str,
+    kind: QueryKind,
     // whether replace '\n' with '\\n', only disable in explain/show create stmts
     replace_newline: bool,
     schema: SchemaRef,
@@ -68,6 +70,7 @@ impl<'a> FormatDisplay<'a> {
             replace_newline,
             schema,
             data,
+            kind: QueryKind::from(query),
             rows: 0,
             progress: None,
             start,
@@ -205,12 +208,20 @@ impl<'a> FormatDisplay<'a> {
         if !self.settings.show_stats {
             return;
         }
+
         if let Some(ref mut stats) = self.stats {
             stats.normalize();
-            let rows_str = if self.rows > 1 { "rows" } else { "row" };
+
+            let (rows, kind) = if matches!(self.kind, QueryKind::Update) {
+                (stats.write_rows, "written")
+            } else {
+                (self.rows, "result")
+            };
+
+            let rows_str = if rows > 1 { "rows" } else { "row" };
             eprintln!(
-                "{} {} in {:.3} sec. Processed {} rows, {} ({} rows/s, {}/s)",
-                self.rows,
+                "{} {} {kind} in {:.3} sec. Processed {} rows, {} ({} rows/s, {}/s)",
+                rows,
                 rows_str,
                 self.start.elapsed().as_secs_f64(),
                 humanize_count(stats.total_rows as f64),
diff --git a/cli/src/session.rs b/cli/src/session.rs
index 46d1b87e1..35a045390 100644
--- a/cli/src/session.rs
+++ b/cli/src/session.rs
@@ -284,8 +284,8 @@ impl Session {
 
         let start = Instant::now();
         let kind = QueryKind::from(query);
-        match kind {
-            QueryKind::Update => {
+        match (kind, is_repl) {
+            (QueryKind::Update, false) => {
                 let affected = self.conn.exec(query).await?;
                 if is_repl {
                     if affected > 0 {
@@ -301,7 +301,7 @@ impl Session {
                 }
                 Ok(false)
             }
-            QueryKind::Query | QueryKind::Explain => {
+            _ => {
                 let replace_newline = replace_newline_in_box_display(query);
                 let (schema, data) = self.conn.query_iter_ext(query).await?;
                 let mut displayer = FormatDisplay::new(
diff --git a/licenserc.toml b/licenserc.toml
index 1aab0bad8..a0c9604d0 100644
--- a/licenserc.toml
+++ b/licenserc.toml
@@ -36,6 +36,9 @@ excludes = [
     "core/tests/**",
     "driver/tests/**",
     "cli/tests/**",
+
+    # python files
+    "**/*.pyi"
 ]
 
 [properties]