fix: improve taos shell user experience for connecting cloud (#20528)

tools/shell/inc/shellInt.h

@@ -28,6 +28,10 @@
 
 #ifdef WEBSOCKET
 #include "taosws.h"
+
+#define SHELL_WS_TIMEOUT                        30
+#define SHELL_WS_DSN_BUFF                       256
+#define SHELL_WS_DSN_MASK                       10
 #endif
 
 #define SHELL_MAX_HISTORY_SIZE                 1000
@@ -99,7 +103,7 @@ typedef struct {
   bool            exit;
 #ifdef WEBSOCKET
   WS_TAOS*        ws_conn;
-  bool		      stop_query;
+  bool            stop_query;
 #endif
 } SShellObj;
 
@@ -139,7 +143,7 @@ void    shellExit();
 void shellTestNetWork();
 
 #ifdef WEBSOCKET
-void	shellCheckConnectMode();
+void shellCheckConnectMode();
 // shellWebsocket.c
 int shell_conn_ws_server(bool first);
 int32_t shell_run_websocket();

tools/shell/src/shellArguments.c

@@ -60,7 +60,7 @@
 #ifdef WEBSOCKET
 #define SHELL_DSN      "The dsn to use when connecting to cloud server."
 #define SHELL_REST     "Use restful mode when connecting."
-#define SHELL_TIMEOUT  "Set the timeout for websocket query in seconds, default is 10."
+#define SHELL_TIMEOUT  "Set the timeout for websocket query in seconds, default is 30."
 #endif
 
 static int32_t shellParseSingleOpt(int32_t key, char *arg);
@@ -127,7 +127,7 @@ static struct argp_option shellOptions[] = {
 #ifdef WEBSOCKET
     {"dsn", 'E', "DSN", 0, SHELL_DSN},
     {"restful", 'R', 0, 0, SHELL_REST},
-	  {"timeout", 'T', "SECONDS", 0, SHELL_TIMEOUT},
+    {"timeout", 'T', "SECONDS", 0, SHELL_TIMEOUT},
 #endif
     {"pktnum", 'N', "PKTNUM", 0, SHELL_PKT_NUM},
     {0},
@@ -223,9 +223,9 @@ static int32_t shellParseSingleOpt(int32_t key, char *arg) {
       pArgs->dsn = arg;
       pArgs->cloud = true;
       break;
-	case 'T':
-	  pArgs->timeout = atoi(arg);
-	  break;
+    case 'T':
+      pArgs->timeout = atoi(arg);
+      break;
 #endif
     case 'V':
       pArgs->is_version = true;
@@ -246,7 +246,8 @@ int32_t shellParseArgsWithoutArgp(int argc, char *argv[]) {
   SShellArgs *pArgs = &shell.args;
 
   for (int i = 1; i < argc; i++) {
-    if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "--usage") == 0 || strcmp(argv[i], "-?") == 0 || strcmp(argv[i], "/?") == 0) {
+    if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "--usage") == 0
+            || strcmp(argv[i], "-?") == 0 || strcmp(argv[i], "/?") == 0) {
       shellParseSingleOpt('?', NULL);
       return 0;
     }
@@ -262,12 +263,14 @@ int32_t shellParseArgsWithoutArgp(int argc, char *argv[]) {
       return -1;
     }
 
-    if (key[1] == 'h' || key[1] == 'P' || key[1] == 'u' || key[1] == 'a' || key[1] == 'c' || key[1] == 's' ||
-        key[1] == 'f' || key[1] == 'd' || key[1] == 'w' || key[1] == 'n' || key[1] == 'l' || key[1] == 'N' 
+    if (key[1] == 'h' || key[1] == 'P' || key[1] == 'u'
+            || key[1] == 'a' || key[1] == 'c' || key[1] == 's'
+            || key[1] == 'f' || key[1] == 'd' || key[1] == 'w'
+            || key[1] == 'n' || key[1] == 'l' || key[1] == 'N'
 #ifdef WEBSOCKET
-	   || key[1] == 'E' || key[1] == 'T'
+        || key[1] == 'E' || key[1] == 'T'
 #endif
-		) {
+    ) {
       if (i + 1 >= argc) {
         fprintf(stderr, "option %s requires an argument\r\n", key);
         return -1;
@@ -279,12 +282,14 @@ int32_t shellParseArgsWithoutArgp(int argc, char *argv[]) {
       }
       shellParseSingleOpt(key[1], val);
       i++;
-    } else if (key[1] == 'p' || key[1] == 'A' || key[1] == 'C' || key[1] == 'r' || key[1] == 'k' || 
-               key[1] == 't' || key[1] == 'V' || key[1] == '?' || key[1] == 1
+    } else if (key[1] == 'p' || key[1] == 'A' || key[1] == 'C'
+                || key[1] == 'r' || key[1] == 'k'
+                || key[1] == 't' || key[1] == 'V'
+                || key[1] == '?' || key[1] == 1
 #ifdef WEBSOCKET
-			   ||key[1] == 'R'
+            ||key[1] == 'R'
 #endif
-			   ) {
+    ) {
       shellParseSingleOpt(key[1], NULL);
     } else {
       fprintf(stderr, "invalid option %s\r\n", key);
@@ -406,7 +411,7 @@ static int32_t shellCheckArgs() {
 
 int32_t shellParseArgs(int32_t argc, char *argv[]) {
   shellInitArgs(argc, argv);
-  shell.info.clientVersion = 
+  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";
   strcpy(shell.info.cusName, cusName);

tools/shell/src/shellEngine.c

@@ -102,10 +102,10 @@ int32_t shellRunSingleCommand(char *command) {
   }
 #ifdef WEBSOCKET
   if (shell.args.restful || shell.args.cloud) {
-	shellRunSingleCommandWebsocketImp(command);
+    shellRunSingleCommandWebsocketImp(command);
   } else {
 #endif
-	shellRunSingleCommandImp(command);
+    shellRunSingleCommandImp(command);
 #ifdef WEBSOCKET
   }
 #endif
@@ -1025,15 +1025,15 @@ void *shellCancelHandler(void *arg) {
     }
 
 #ifdef WEBSOCKET
-	if (shell.args.restful || shell.args.cloud) {
-		shell.stop_query = true;
-	} else {
+    if (shell.args.restful || shell.args.cloud) {
+      shell.stop_query = true;
+    } else {
 #endif
-    if (shell.conn) {
-		  taos_kill_query(shell.conn);
-		}
+      if (shell.conn) {
+        taos_kill_query(shell.conn);
+      }
 #ifdef WEBSOCKET
-	}
+    }
 #endif
   #ifdef WINDOWS
     printf("\n%s", shell.info.promptHeader);
@@ -1083,21 +1083,21 @@ int32_t shellExecute() {
   SShellArgs *pArgs = &shell.args;
 #ifdef WEBSOCKET
   if (shell.args.restful || shell.args.cloud) {
-	if (shell_conn_ws_server(1)) {
-		return -1;
-	}
+    if (shell_conn_ws_server(1)) {
+      return -1;
+    }
   } else {
 #endif
-	if (shell.args.auth == NULL) {
-		shell.conn = taos_connect(pArgs->host, pArgs->user, pArgs->password, pArgs->database, pArgs->port);
-	} else {
-		shell.conn = taos_connect_auth(pArgs->host, pArgs->user, pArgs->auth, pArgs->database, pArgs->port);
-	}
-
-	if (shell.conn == NULL) {
-		fflush(stdout);
-		return -1;
-	}
+    if (shell.args.auth == NULL) {
+      shell.conn = taos_connect(pArgs->host, pArgs->user, pArgs->password, pArgs->database, pArgs->port);
+    } else {
+      shell.conn = taos_connect_auth(pArgs->host, pArgs->user, pArgs->auth, pArgs->database, pArgs->port);
+    }
+
+    if (shell.conn == NULL) {
+      fflush(stdout);
+      return -1;
+    }
 #ifdef WEBSOCKET
   }
 #endif
@@ -1118,13 +1118,13 @@ int32_t shellExecute() {
       shellSourceFile(pArgs->file);
     }
 #ifdef WEBSOCKET
-	if (shell.args.restful || shell.args.cloud) {
-		ws_close(shell.ws_conn);
-	} else {
+    if (shell.args.restful || shell.args.cloud) {
+      ws_close(shell.ws_conn);
+    } else {
 #endif
-		taos_close(shell.conn);
+      taos_close(shell.conn);
 #ifdef WEBSOCKET
-	}
+    }
 #endif
 
     shellWriteHistory();
@@ -1148,9 +1148,9 @@ int32_t shellExecute() {
   if (!shell.args.restful && !shell.args.cloud) {
 #endif
 #ifndef WINDOWS
-  printfIntroduction();
+    printfIntroduction();
 #endif  
-	shellGetGrantInfo();
+    shellGetGrantInfo();
 #ifdef WEBSOCKET
   }
 #endif

tools/shell/src/shellMain.c

@@ -45,7 +45,7 @@ void shellCrashHandler(int signum, void *sigInfo, void *context) {
 int main(int argc, char *argv[]) {
   shell.exit = false;
 #ifdef WEBSOCKET
-  shell.args.timeout = 10;
+  shell.args.timeout = SHELL_WS_TIMEOUT;
   shell.args.cloud = true;
 #endif
 

tools/shell/src/shellWebsocket.c

@@ -14,22 +14,51 @@
  * along with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 #ifdef WEBSOCKET
-#include "taosws.h"
-#include "shellInt.h"
+#include <taosws.h>
+#include <shellInt.h>
 
 int shell_conn_ws_server(bool first) {
-  shell.ws_conn = ws_connect_with_dsn(shell.args.dsn);
-  if (!shell.ws_conn) {
-    fprintf(stderr, "failed to connect %s, reason: %s\n",
-        shell.args.dsn, ws_errstr(NULL));
+  char cuttedDsn[SHELL_WS_DSN_BUFF] = {0};
+  int dsnLen = strlen(shell.args.dsn);
+  snprintf(cuttedDsn,
+           ((dsnLen-SHELL_WS_DSN_MASK) > SHELL_WS_DSN_BUFF)?
+            SHELL_WS_DSN_BUFF:(dsnLen-SHELL_WS_DSN_MASK),
+           "%s", shell.args.dsn);
+  fprintf(stdout, "trying to connect %s*** ", cuttedDsn);
+  fflush(stdout);
+  for (int i = 0; i < shell.args.timeout; i++) {
+    shell.ws_conn = ws_connect_with_dsn(shell.args.dsn);
+    if (NULL == shell.ws_conn) {
+      int errNo = ws_errno(NULL);
+      if (0xE001 == errNo) {
+        fprintf(stdout, ".");
+        fflush(stdout);
+        taosMsleep(1000);  // sleep 1 second then try again
+        continue;
+      } else {
+        fprintf(stderr, "\nfailed to connect %s***, reason: %s\n",
+            cuttedDsn, ws_errstr(NULL));
+        return -1;
+      }
+    } else {
+      break;
+    }
+  }
+  if (NULL == shell.ws_conn) {
+    fprintf(stdout, "\n timeout\n");
+    fprintf(stderr, "\nfailed to connect %s***, reason: %s\n",
+      cuttedDsn, ws_errstr(NULL));
     return -1;
+  } else {
+    fprintf(stdout, "\n");
   }
   if (first && shell.args.restful) {
-    fprintf(stdout, "successfully connect to %s\n\n",
+    fprintf(stdout, "successfully connected to %s\n\n",
         shell.args.dsn);
   } else if (first && shell.args.cloud) {
-    fprintf(stdout, "successfully connect to cloud service\n");
+    fprintf(stdout, "successfully connected to cloud service\n");
   }
+  fflush(stdout);
   return 0;
 }
 
@@ -118,7 +147,8 @@ static int verticalPrintWebsocket(WS_RES* wres, double* pexecute_time) {
   return numOfRows;
 }
 
-static int dumpWebsocketToFile(const char* fname, WS_RES* wres, double* pexecute_time) {
+static int dumpWebsocketToFile(const char* fname, WS_RES* wres,
+                               double* pexecute_time) {
   char fullname[PATH_MAX] = {0};
   if (taosExpandDir(fname, fullname, PATH_MAX) != 0) {
     tstrncpy(fullname, fname, PATH_MAX);
@@ -150,7 +180,7 @@ static int dumpWebsocketToFile(const char* fname, WS_RES* wres, double* pexecute
     }
     taosFprintfFile(pFile, "%s", fields[col].name);
   }
-  taosFprintfFile(pFile, "\r\n"); 
+  taosFprintfFile(pFile, "\r\n");
   do {
     uint8_t ty;
     uint32_t len;
@@ -161,7 +191,8 @@ static int dumpWebsocketToFile(const char* fname, WS_RES* wres, double* pexecute
           taosFprintfFile(pFile, ",");
         }
         const void *value = ws_get_value_in_block(wres, i, j, &ty, &len);
-        shellDumpFieldToFile(pFile, (const char*)value, fields + j, len, precision);
+        shellDumpFieldToFile(pFile, (const char*)value,
+                             fields + j, len, precision);
       }
       taosFprintfFile(pFile, "\r\n");
     }
@@ -171,7 +202,9 @@ static int dumpWebsocketToFile(const char* fname, WS_RES* wres, double* pexecute
   return numOfRows;
 }
 
-static int shellDumpWebsocket(WS_RES *wres, char *fname, int *error_no, bool vertical, double* pexecute_time) {
+static int shellDumpWebsocket(WS_RES *wres, char *fname,
+                              int *error_no, bool vertical,
+                              double* pexecute_time) {
   int numOfRows = 0;
   if (fname != NULL) {
     numOfRows = dumpWebsocketToFile(fname, wres, pexecute_time);
@@ -227,13 +260,16 @@ void shellRunSingleCommandWebsocketImp(char *command) {
       // if it's not a ws connection error
       if (TSDB_CODE_WS_DSN_ERROR != (code&TSDB_CODE_WS_DSN_ERROR)) {
         et = taosGetTimestampUs();
-        fprintf(stderr, "\nDB: error: %s (%.6fs)\n", ws_errstr(res), (et - st)/1E6);
+        fprintf(stderr, "\nDB: error: %s (%.6fs)\n",
+                ws_errstr(res), (et - st)/1E6);
         ws_free_result(res);
         return;
       }
-      if (code == TSDB_CODE_WS_SEND_TIMEOUT || code == TSDB_CODE_WS_RECV_TIMEOUT) {
+      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) {
+      } else if (code == TSDB_CODE_WS_INTERNAL_ERRO
+                    || code == TSDB_CODE_WS_CLOSED) {
         shell.ws_conn = NULL;
       }
       ws_free_result(res);
@@ -252,7 +288,8 @@ void shellRunSingleCommandWebsocketImp(char *command) {
     execute_time = ws_take_timing(res)/1E6;
   }
 
-  if (shellRegexMatch(command, "^\\s*use\\s+[a-zA-Z0-9_]+\\s*;\\s*$", REG_EXTENDED | REG_ICASE)) {
+  if (shellRegexMatch(command, "^\\s*use\\s+[a-zA-Z0-9_]+\\s*;\\s*$",
+                      REG_EXTENDED | REG_ICASE)) {
     fprintf(stdout, "Database changed.\r\n\r\n");
     fflush(stdout);
     ws_free_result(res);
@@ -266,10 +303,12 @@ void shellRunSingleCommandWebsocketImp(char *command) {
     double total_time = (et - st)/1E3;
     double net_time = total_time - (double)execute_time;
     printf("Query Ok, %d of %d row(s) in database\n", numOfRows, numOfRows);
-    printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n", execute_time, net_time, total_time);
+    printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n",
+           execute_time, net_time, total_time);
   } else {
     int error_no = 0;
-    numOfRows  = shellDumpWebsocket(res, fname, &error_no, printMode, &execute_time);
+    numOfRows  = shellDumpWebsocket(res, fname, &error_no,
+                                    printMode, &execute_time);
     if (numOfRows < 0) {
       ws_free_result(res);
       return;
@@ -279,11 +318,13 @@ void shellRunSingleCommandWebsocketImp(char *command) {
     double net_time = total_time - execute_time;
     if (error_no == 0 && !shell.stop_query) {
       printf("Query OK, %d row(s) in set\n", numOfRows);
-      printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n", execute_time, net_time, total_time);
+      printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n",
+             execute_time, net_time, total_time);
     } else {
       printf("Query interrupted, %d row(s) in set (%.6fs)\n", numOfRows,
           (et - st)/1E6);
-      printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n", execute_time, net_time, total_time);
+      printf("Execute: %.2f ms Network: %.2f ms Total: %.2f ms\n",
+             execute_time, net_time, total_time);
     }
   }
   printf("\n");