From 368c6d996b73531d56b9c24b542258dc5bf1d088 Mon Sep 17 00:00:00 2001
From: Alex Duan <417921451@qq.com>
Date: Wed, 25 Oct 2023 16:46:32 +0800
Subject: [PATCH] feat: add env TDENGINE_DSN to taos-CLI

---
 tools/shell/inc/shellInt.h       |  1 +
 tools/shell/src/shellArguments.c |  2 +-
 tools/shell/src/shellMain.c      |  1 +
 tools/shell/src/shellUtil.c      | 12 +++++++++++-
 tools/shell/src/shellWebsocket.c | 12 +++++++++++-
 5 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/tools/shell/inc/shellInt.h b/tools/shell/inc/shellInt.h
index 57415f8335..1c885b151c 100644
--- a/tools/shell/inc/shellInt.h
+++ b/tools/shell/inc/shellInt.h
@@ -80,6 +80,7 @@ typedef struct {
 #ifdef WEBSOCKET
   bool        restful;
   bool        cloud;
+  bool        local;
   char*       dsn;
   int32_t     timeout;
 #endif
diff --git a/tools/shell/src/shellArguments.c b/tools/shell/src/shellArguments.c
index 52cb524e3b..4817b23029 100644
--- a/tools/shell/src/shellArguments.c
+++ b/tools/shell/src/shellArguments.c
@@ -410,7 +410,7 @@ int32_t shellParseArgs(int32_t argc, char *argv[]) {
   shellInitArgs(argc, argv);
   shell.info.clientVersion =
       "Welcome to the %s Command Line Interface, Client Version:%s\r\n"
-      "Copyright (c) 2022 by %s, all rights reserved.\r\n\r\n";
+      "Copyright (c) 2023 by %s, all rights reserved.\r\n\r\n";
 #ifdef CUS_NAME
   strcpy(shell.info.cusName, CUS_NAME);
 #else
diff --git a/tools/shell/src/shellMain.c b/tools/shell/src/shellMain.c
index 795621dfdd..18f4ca21d1 100644
--- a/tools/shell/src/shellMain.c
+++ b/tools/shell/src/shellMain.c
@@ -47,6 +47,7 @@ int main(int argc, char *argv[]) {
 #ifdef WEBSOCKET
   shell.args.timeout = SHELL_WS_TIMEOUT;
   shell.args.cloud = true;
+  shell.args.local = false;
 #endif
 
 #if !defined(WINDOWS)
diff --git a/tools/shell/src/shellUtil.c b/tools/shell/src/shellUtil.c
index 93451c85a9..f57c449961 100644
--- a/tools/shell/src/shellUtil.c
+++ b/tools/shell/src/shellUtil.c
@@ -130,11 +130,21 @@ void shellCheckConnectMode() {
 	}
 	if (shell.args.cloud) {
 		shell.args.dsn = getenv("TDENGINE_CLOUD_DSN");
-		if (shell.args.dsn) {
+		if (shell.args.cloud_dsn) {
 			shell.args.cloud = true;
+      shell.args.local = false;
 			shell.args.restful = false;
 			return;
 		}
+
+    shell.args.dsn = getenv("TDENGINE_DSN");
+		if (shell.args.dsn) {
+			shell.args.cloud = true;
+      shell.args.local = true;
+			shell.args.restful = false;
+			return;
+		}
+
 		if (shell.args.restful) {
 			if (!shell.args.host) {
 				shell.args.host = "localhost";
diff --git a/tools/shell/src/shellWebsocket.c b/tools/shell/src/shellWebsocket.c
index e83ceff099..15affd3c76 100644
--- a/tools/shell/src/shellWebsocket.c
+++ b/tools/shell/src/shellWebsocket.c
@@ -59,7 +59,17 @@ int shell_conn_ws_server(bool first) {
     fprintf(stdout, "successfully connected to %s\n\n",
         shell.args.dsn);
   } else if (first && shell.args.cloud) {
-    fprintf(stdout, "successfully connected to cloud service\n");
+    if(shell.args.local) {
+      const char* host = strstr(shell.args.dsn, "@");
+      if(host) {
+        host += 1;
+      } else {
+        host = shell.args.dsn;
+      }
+      fprintf(stdout, "successfully connected to %s\n", host);
+    } else {
+      fprintf(stdout, "successfully connected to cloud service\n");
+    }
   }
   fflush(stdout);