From d3a21816ccefd85a616454c18d413b6d5c936708 Mon Sep 17 00:00:00 2001
From: afwerar <1296468573@qq.com>
Date: Wed, 21 Sep 2022 11:43:36 +0800
Subject: [PATCH] shell: taos shell websocket reconnect error

---
 tools/shell/src/shellUtil.c      |  2 +-
 tools/shell/src/shellWebsocket.c | 39 ++++++++++++++++++--------------
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/tools/shell/src/shellUtil.c b/tools/shell/src/shellUtil.c
index 0430428c38..8c47d16555 100644
--- a/tools/shell/src/shellUtil.c
+++ b/tools/shell/src/shellUtil.c
@@ -143,7 +143,7 @@ void shellCheckConnectMode() {
 				shell.args.port = 6041;
 			}
 			shell.args.dsn = taosMemoryCalloc(1, 1024);
-			snprintf(shell.args.dsn, 1024, "ws://%s:%d/rest/ws",
+			snprintf(shell.args.dsn, 1024, "ws://%s:%d",
 					shell.args.host, shell.args.port);
 		}
 		shell.args.cloud = false;
diff --git a/tools/shell/src/shellWebsocket.c b/tools/shell/src/shellWebsocket.c
index b8b8392b96..94bb909e29 100644
--- a/tools/shell/src/shellWebsocket.c
+++ b/tools/shell/src/shellWebsocket.c
@@ -206,26 +206,31 @@ void shellRunSingleCommandWebsocketImp(char *command) {
     printMode = true;  // When output to a file, the switch does not work.
   }
 
-  if (!shell.ws_conn && shell_conn_ws_server(0)) {
-    return;
-  }
-
   shell.stop_query = false;
-  st = taosGetTimestampUs();
+  WS_RES* res;
 
-  WS_RES* res = ws_query_timeout(shell.ws_conn, command, shell.args.timeout);
-  int code = ws_errno(res);
-  if (code != 0) {
-    et = taosGetTimestampUs();
-    fprintf(stderr, "\nDB: error: %s (%.6fs)\n", ws_errstr(res), (et - st)/1E6);
-    if (code == TSDB_CODE_WS_SEND_TIMEOUT || code == TSDB_CODE_WS_RECV_TIMEOUT) {
-      fprintf(stderr, "Hint: use -t to increase the timeout in seconds\n");
-    } else if (code == TSDB_CODE_WS_INTERNAL_ERRO || code == TSDB_CODE_WS_CLOSED) {
-      fprintf(stderr, "TDengine server is down, will try to reconnect\n");
-      shell.ws_conn = NULL;
+  for (int reconnectNum = 0; reconnectNum < 2; reconnectNum++) {
+    if (!shell.ws_conn && shell_conn_ws_server(0)) {
+      return;
     }
-    ws_free_result(res);
-    return;
+    st = taosGetTimestampUs();
+
+    res = ws_query_timeout(shell.ws_conn, command, shell.args.timeout);
+    int code = ws_errno(res);
+    if (code != 0 && !shell.stop_query) {
+      et = taosGetTimestampUs();
+      fprintf(stderr, "\nDB: error: %s (%.6fs)\n", ws_errstr(res), (et - st)/1E6);
+      if (code == TSDB_CODE_WS_SEND_TIMEOUT || code == TSDB_CODE_WS_RECV_TIMEOUT) {
+        fprintf(stderr, "Hint: use -t to increase the timeout in seconds\n");
+      } else if (code == TSDB_CODE_WS_INTERNAL_ERRO || code == TSDB_CODE_WS_CLOSED) {
+        fprintf(stderr, "TDengine server is down, will try to reconnect\n");
+        shell.ws_conn = NULL;
+      }
+      ws_free_result(res);
+      if (reconnectNum == 0) continue;
+      return;
+    }
+    break;
   }
 
   double execute_time = ws_take_timing(res)/1E6;