From 6f15e4cd2a9292465f400e729978d4997d3b85ce Mon Sep 17 00:00:00 2001
From: sundyli <543950155@qq.com>
Date: Thu, 6 Apr 2023 08:50:48 -0700
Subject: [PATCH] feat: cli support reconnect (#34)

---
 cli/Cargo.toml     |  2 +-
 cli/src/session.rs | 59 +++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index ee7ebd1fb..24c3132b2 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "bendsql"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 license = "Apache-2.0"
 description = "Databend Native Cli tool"
diff --git a/cli/src/session.rs b/cli/src/session.rs
index 66779aa67..e0abf8647 100644
--- a/cli/src/session.rs
+++ b/cli/src/session.rs
@@ -35,6 +35,10 @@ use crate::helper::CliHelper;
 pub struct Session {
     client: FlightSqlServiceClient,
     is_repl: bool,
+    endpoint: Endpoint,
+    user: String,
+    password: String,
+
     config: Config,
     prompt: String,
 }
@@ -47,6 +51,12 @@ impl Session {
         password: &str,
         is_repl: bool,
     ) -> Result<Self, ArrowError> {
+        if is_repl {
+            println!("Welcome to databend-cli.");
+            println!("Connecting to {} as user {}.", endpoint.uri(), user);
+            println!();
+        }
+
         let channel = endpoint
             .connect()
             .await
@@ -54,12 +64,6 @@ impl Session {
 
         let mut client = FlightSqlServiceClient::new(channel);
 
-        if is_repl {
-            println!("Welcome to databend-cli.");
-            println!("Connecting to {} as user {}.", endpoint.uri(), user);
-            println!();
-        }
-
         // enable progress
         client.set_header("bendsql", "1");
         let _token = client.handshake(user, password).await.unwrap();
@@ -76,6 +80,9 @@ impl Session {
             client,
             is_repl,
             prompt,
+            user: user.to_owned(),
+            password: password.to_owned(),
+            endpoint,
         })
     }
 
@@ -127,9 +134,16 @@ impl Session {
                         break;
                     }
                     Ok(false) => {}
-                    Err(e) => {
-                        eprintln!("{}", format_error(e));
-                    }
+                    Err(e) => match e {
+                        ArrowError::IoError(e) if e.contains("Unauthenticated") => {
+                            if let Err(e) = self.reconnect().await {
+                                eprintln!("Reconnect error: {}", format_error(e));
+                            } else if let Err(e) = self.handle_query(true, &query).await {
+                                eprintln!("{}", format_error(e));
+                            }
+                        }
+                        e => eprintln!("{}", format_error(e)),
+                    },
                 }
             }
             query.clear();
@@ -200,6 +214,33 @@ impl Session {
 
         Ok(false)
     }
+
+    async fn reconnect(&mut self) -> Result<(), ArrowError> {
+        if self.is_repl {
+            println!(
+                "Connecting to {} as user {}.",
+                self.endpoint.uri(),
+                self.user
+            );
+            println!();
+        }
+
+        let channel = self
+            .endpoint
+            .connect()
+            .await
+            .map_err(|err| ArrowError::IoError(err.to_string()))?;
+
+        self.client = FlightSqlServiceClient::new(channel);
+        // enable progress
+        self.client.set_header("bendsql", "1");
+        let _token = self
+            .client
+            .handshake(&self.user, &self.password)
+            .await
+            .unwrap();
+        Ok(())
+    }
 }
 
 fn get_history_path() -> String {