From ff665c362f6ef3740133314c8b36106bdca8feac Mon Sep 17 00:00:00 2001
From: Amos Bird <amosbird@gmail.com>
Date: Sun, 11 Oct 2020 10:09:08 +0800
Subject: [PATCH] Add print query id support to client

---
 programs/client/Client.cpp                    | 37 ++++++++++++++++---
 src/Interpreters/Context.cpp                  |  3 --
 .../01520_client_print_query_id.reference     |  0
 .../01520_client_print_query_id.sh            | 22 +++++++++++
 4 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 tests/queries/0_stateless/01520_client_print_query_id.reference
 create mode 100755 tests/queries/0_stateless/01520_client_print_query_id.sh

diff --git a/programs/client/Client.cpp b/programs/client/Client.cpp
index df365a2089e..ea46b25e4bc 100644
--- a/programs/client/Client.cpp
+++ b/programs/client/Client.cpp
@@ -134,7 +134,6 @@ private:
     bool stdout_is_a_tty = false;        /// stdout is a terminal.
 
     std::unique_ptr<Connection> connection;    /// Connection to DB.
-    String query_id;                     /// Current query_id.
     String full_query; /// Current query as it was given to the client.
 
     // Current query as it will be sent to the server. It may differ from the
@@ -219,6 +218,8 @@ private:
     QueryFuzzer fuzzer;
     int query_fuzzer_runs = 0;
 
+    std::vector<std::pair<String, String>> query_id_formats;
+
     void initialize(Poco::Util::Application & self) override
     {
         Poco::Util::Application::initialize(self);
@@ -243,6 +244,17 @@ private:
         /// Set path for format schema files
         if (config().has("format_schema_path"))
             context.setFormatSchemaPath(Poco::Path(config().getString("format_schema_path")).toString());
+
+        /// Initialize query_id_formats if any
+        if (config().has("query_id_formats"))
+        {
+            Poco::Util::AbstractConfiguration::Keys keys;
+            config().keys("query_id_formats", keys);
+            for (const auto & name : keys)
+                query_id_formats.emplace_back(name + ":", config().getString("query_id_formats." + name));
+        }
+        if (query_id_formats.empty())
+            query_id_formats.emplace_back("Query id:", " {query_id}");
     }
 
 
@@ -559,7 +571,7 @@ private:
 
         if (is_interactive)
         {
-            if (!query_id.empty())
+            if (config().has("query_id"))
                 throw Exception("query_id could be specified only in non-interactive mode", ErrorCodes::BAD_ARGUMENTS);
             if (print_time_to_stderr)
                 throw Exception("time option could be specified only in non-interactive mode", ErrorCodes::BAD_ARGUMENTS);
@@ -665,7 +677,9 @@ private:
         }
         else
         {
-            query_id = config().getString("query_id", "");
+            auto query_id = config().getString("query_id", "");
+            if (!query_id.empty())
+                context.setCurrentQueryId(query_id);
             if (query_fuzzer_runs)
             {
                 nonInteractiveWithFuzzing();
@@ -1274,6 +1288,19 @@ private:
             std_out.next();
         }
 
+        if (is_interactive)
+        {
+            // Generate a new query_id
+            context.setCurrentQueryId("");
+            for (const auto & query_id_format : query_id_formats)
+            {
+                writeString(query_id_format.first, std_out);
+                writeString(fmt::format(query_id_format.second, fmt::arg("query_id", context.getCurrentQueryId())), std_out);
+                writeChar('\n', std_out);
+                std_out.next();
+            }
+        }
+
         watch.restart();
         processed_rows = 0;
         progress.reset();
@@ -1399,7 +1426,7 @@ private:
                 connection->sendQuery(
                     connection_parameters.timeouts,
                     query_to_send,
-                    query_id,
+                    context.getCurrentQueryId(),
                     QueryProcessingStage::Complete,
                     &context.getSettingsRef(),
                     &context.getClientInfo(),
@@ -1440,7 +1467,7 @@ private:
         connection->sendQuery(
             connection_parameters.timeouts,
             query_to_send,
-            query_id,
+            context.getCurrentQueryId(),
             QueryProcessingStage::Complete,
             &context.getSettingsRef(),
             &context.getClientInfo(),
diff --git a/src/Interpreters/Context.cpp b/src/Interpreters/Context.cpp
index 68f18ffffe8..b2f2b9c83ce 100644
--- a/src/Interpreters/Context.cpp
+++ b/src/Interpreters/Context.cpp
@@ -1111,9 +1111,6 @@ void Context::setCurrentDatabase(const String & name)
 
 void Context::setCurrentQueryId(const String & query_id)
 {
-    if (!client_info.current_query_id.empty())
-        throw Exception("Logical error: attempt to set query_id twice", ErrorCodes::LOGICAL_ERROR);
-
     String query_id_to_set = query_id;
 
     if (query_id_to_set.empty())    /// If the user did not submit his query_id, then we generate it ourselves.
diff --git a/tests/queries/0_stateless/01520_client_print_query_id.reference b/tests/queries/0_stateless/01520_client_print_query_id.reference
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01520_client_print_query_id.sh b/tests/queries/0_stateless/01520_client_print_query_id.sh
new file mode 100755
index 00000000000..4c70aeb966e
--- /dev/null
+++ b/tests/queries/0_stateless/01520_client_print_query_id.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/expect -f
+
+log_user 0
+set timeout 5
+match_max 100000
+
+if ![info exists env(CLICKHOUSE_PORT_TCP)] {set env(CLICKHOUSE_PORT_TCP) 9000}
+
+spawn clickhouse-client --port "$env(CLICKHOUSE_PORT_TCP)"
+expect ":) "
+
+# Make a query
+send -- "SELECT 'print query id'\r"
+expect {
+    "Query id: *" { }
+    timeout { exit 1 }
+}
+expect "print query id"
+expect ":) "
+
+send -- "\4"
+expect eof