From e3653e70da7cdef26ae7b73b5f7544502890be97 Mon Sep 17 00:00:00 2001 From: Alexey Milovidov Date: Tue, 15 Jun 2021 04:43:35 +0300 Subject: [PATCH] Update prompt in client when reconnecting --- programs/client/Client.cpp | 126 +++++++++++++++++++------------------ src/Client/Connection.h | 3 + 2 files changed, 68 insertions(+), 61 deletions(-) diff --git a/programs/client/Client.cpp b/programs/client/Client.cpp index 49d3dc3d10b..1d0d789a29a 100644 --- a/programs/client/Client.cpp +++ b/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 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 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) diff --git a/src/Client/Connection.h b/src/Client/Connection.h index 80dbd9ed44e..2ea5a236a13 100644 --- a/src/Client/Connection.h +++ b/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);