Merge pull request #25281 from ClickHouse/client-update-prompt-on-reconnect

Update prompt in client when reconnecting

programs/client/Client.cpp

@@ -549,65 +549,6 @@ private:
 
         /// Initialize DateLUT here to avoid counting time spent here as query execution time.
         const auto local_tz = DateLUT::instance().getTimeZone();
-        if (!context->getSettingsRef().use_client_time_zone)
-        {
-            const auto & time_zone = connection->getServerTimezone(connection_parameters.timeouts);
-            if (!time_zone.empty())
-            {
-                try
-                {
-                    DateLUT::setDefaultTimezone(time_zone);
-                }
-                catch (...)
-                {
-                    std::cerr << "Warning: could not switch to server time zone: " << time_zone
-                              << ", reason: " << getCurrentExceptionMessage(/* with_stacktrace = */ false) << std::endl
-                              << "Proceeding with local time zone." << std::endl
-                              << std::endl;
-                }
-            }
-            else
-            {
-                std::cerr << "Warning: could not determine server time zone. "
-                          << "Proceeding with local time zone." << std::endl
-                          << std::endl;
-            }
-        }
-
-        Strings keys;
-
-        prompt_by_server_display_name = config().getRawString("prompt_by_server_display_name.default", "{display_name} :) ");
-
-        config().keys("prompt_by_server_display_name", keys);
-
-        for (const String & key : keys)
-        {
-            if (key != "default" && server_display_name.find(key) != std::string::npos)
-            {
-                prompt_by_server_display_name = config().getRawString("prompt_by_server_display_name." + key);
-                break;
-            }
-        }
-
-        /// Prompt may contain escape sequences including \e[ or \x1b[ sequences to set terminal color.
-        {
-            String unescaped_prompt_by_server_display_name;
-            ReadBufferFromString in(prompt_by_server_display_name);
-            readEscapedString(unescaped_prompt_by_server_display_name, in);
-            prompt_by_server_display_name = std::move(unescaped_prompt_by_server_display_name);
-        }
-
-        /// Prompt may contain the following substitutions in a form of {name}.
-        std::map<String, String> prompt_substitutions{
-            {"host", connection_parameters.host},
-            {"port", toString(connection_parameters.port)},
-            {"user", connection_parameters.user},
-            {"display_name", server_display_name},
-        };
-
-        /// Quite suboptimal.
-        for (const auto & [key, value] : prompt_substitutions)
-            boost::replace_all(prompt_by_server_display_name, "{" + key + "}", value);
 
         if (is_interactive)
         {
@@ -805,6 +746,66 @@ private:
                           << std::endl;
             }
         }
+
+        if (!context->getSettingsRef().use_client_time_zone)
+        {
+            const auto & time_zone = connection->getServerTimezone(connection_parameters.timeouts);
+            if (!time_zone.empty())
+            {
+                try
+                {
+                    DateLUT::setDefaultTimezone(time_zone);
+                }
+                catch (...)
+                {
+                    std::cerr << "Warning: could not switch to server time zone: " << time_zone
+                              << ", reason: " << getCurrentExceptionMessage(/* with_stacktrace = */ false) << std::endl
+                              << "Proceeding with local time zone." << std::endl
+                              << std::endl;
+                }
+            }
+            else
+            {
+                std::cerr << "Warning: could not determine server time zone. "
+                          << "Proceeding with local time zone." << std::endl
+                          << std::endl;
+            }
+        }
+
+        Strings keys;
+
+        prompt_by_server_display_name = config().getRawString("prompt_by_server_display_name.default", "{display_name} :) ");
+
+        config().keys("prompt_by_server_display_name", keys);
+
+        for (const String & key : keys)
+        {
+            if (key != "default" && server_display_name.find(key) != std::string::npos)
+            {
+                prompt_by_server_display_name = config().getRawString("prompt_by_server_display_name." + key);
+                break;
+            }
+        }
+
+        /// Prompt may contain escape sequences including \e[ or \x1b[ sequences to set terminal color.
+        {
+            String unescaped_prompt_by_server_display_name;
+            ReadBufferFromString in(prompt_by_server_display_name);
+            readEscapedString(unescaped_prompt_by_server_display_name, in);
+            prompt_by_server_display_name = std::move(unescaped_prompt_by_server_display_name);
+        }
+
+        /// Prompt may contain the following substitutions in a form of {name}.
+        std::map<String, String> prompt_substitutions{
+            {"host", connection_parameters.host},
+            {"port", toString(connection_parameters.port)},
+            {"user", connection_parameters.user},
+            {"display_name", server_display_name},
+        };
+
+        /// Quite suboptimal.
+        for (const auto & [key, value] : prompt_substitutions)
+            boost::replace_all(prompt_by_server_display_name, "{" + key + "}", value);
     }
 
 
@@ -1202,7 +1203,9 @@ private:
                 client_exception.reset();
                 server_exception.reset();
                 have_error = false;
-                connection->forceConnected(connection_parameters.timeouts);
+
+                if (!connection->checkConnected())
+                    connect();
             }
 
             // Report error.
@@ -1603,7 +1606,8 @@ private:
             if (with_output && with_output->settings_ast)
                 apply_query_settings(*with_output->settings_ast);
 
-            connection->forceConnected(connection_parameters.timeouts);
+            if (!connection->checkConnected())
+                connect();
 
             ASTPtr input_function;
             if (insert && insert->select)

src/Client/Connection.h

@@ -185,6 +185,9 @@ public:
 
     bool isConnected() const { return connected; }
 
+    /// Check if connection is still active with ping request.
+    bool checkConnected() { return connected && ping(); }
+
     TablesStatusResponse getTablesStatus(const ConnectionTimeouts & timeouts,
                                          const TablesStatusRequest & request);