Add ability to pass QueryKind via clickhouse-client/local (useful for debugging)

v2: fix LocalConnection::sendQuery() for Suggest (comes w/o client_info) [1]
    [1]: https://s3.amazonaws.com/clickhouse-test-reports/37290/7c85175963226ff78eec542efafcff4e650aa0f0/stateless_tests__ubsan__actions_.html
Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>

programs/bash-completion/completions/clickhouse-bootstrap

@@ -34,6 +34,12 @@ CLICKHOUSE_QueryProcessingStage=(
     with_mergeable_state_after_aggregation_and_limit
 )
 
+CLICKHOUSE_QueryKind=(
+    initial_query
+    secondary_query
+    no_query
+)
+
 CLICKHOUSE_Format=(
     CapnProto
     PostgreSQLWire
@@ -124,6 +130,10 @@ function _complete_for_clickhouse_generic_bin_impl()
             COMPREPLY=( $(compgen -W "${CLICKHOUSE_QueryProcessingStage[*]}" -- "$cur") )
             return 1
             ;;
+        --query_kind)
+            COMPREPLY=( $(compgen -W "${CLICKHOUSE_QueryKind[*]}" -- "$cur") )
+            return 1
+            ;;
         --send_logs_level)
             COMPREPLY=( $(compgen -W "${CLICKHOUSE_logs_level[*]}" -- "$cur") )
             return 1

programs/client/Client.cpp

@@ -1038,6 +1038,7 @@ void Client::processConfig()
     ClientInfo & client_info = global_context->getClientInfo();
     client_info.setInitialQuery();
     client_info.quota_key = config().getString("quota_key", "");
+    client_info.query_kind = query_kind;
 }
 
 

programs/local/LocalServer.cpp

@@ -626,6 +626,7 @@ void LocalServer::processConfig()
 
     ClientInfo & client_info = global_context->getClientInfo();
     client_info.setInitialQuery();
+    client_info.query_kind = query_kind;
 }
 
 

src/Client/ClientBase.cpp

@@ -119,6 +119,17 @@ namespace ProfileEvents
 namespace DB
 {
 
+static ClientInfo::QueryKind parseQueryKind(const String & query_kind)
+{
+    if (query_kind == "initial_query")
+        return ClientInfo::QueryKind::INITIAL_QUERY;
+    if (query_kind == "secondary_query")
+        return ClientInfo::QueryKind::SECONDARY_QUERY;
+    if (query_kind == "no_query")
+        return ClientInfo::QueryKind::NO_QUERY;
+    throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unknown query kind {}", query_kind);
+}
+
 static void incrementProfileEventsBlock(Block & dst, const Block & src)
 {
     if (!dst)
@@ -2125,6 +2136,7 @@ void ClientBase::init(int argc, char ** argv)
 
         ("query,q", po::value<std::string>(), "query")
         ("stage", po::value<std::string>()->default_value("complete"), "Request query processing up to specified stage: complete,fetch_columns,with_mergeable_state,with_mergeable_state_after_aggregation,with_mergeable_state_after_aggregation_and_limit")
+        ("query_kind", po::value<std::string>()->default_value("initial_query"), "One of initial_query/secondary_query/no_query")
         ("query_id", po::value<std::string>(), "query_id")
         ("progress", "print progress of queries execution")
 
@@ -2255,6 +2267,7 @@ void ClientBase::init(int argc, char ** argv)
         server_logs_file = options["server_logs_file"].as<std::string>();
 
     query_processing_stage = QueryProcessingStage::fromString(options["stage"].as<std::string>());
+    query_kind = parseQueryKind(options["query_kind"].as<std::string>());
     profile_events.print = options.count("print-profile-events");
     profile_events.delay_ms = options["profile-events-delay-ms"].as<UInt64>();
 

src/Client/ClientBase.h

@@ -256,6 +256,7 @@ protected:
     } profile_events;
 
     QueryProcessingStage::Enum query_processing_stage;
+    ClientInfo::QueryKind query_kind;
 
     bool fake_drop = false;
 

src/Client/LocalConnection.cpp

@@ -73,10 +73,14 @@ void LocalConnection::sendQuery(
     const String & query_id,
     UInt64 stage,
     const Settings *,
-    const ClientInfo *,
+    const ClientInfo * client_info,
     bool,
     std::function<void(const Progress &)> process_progress_callback)
 {
+    /// Suggestion comes without client_info.
+    if (client_info)
+        query_context = session.makeQueryContext(*client_info);
+    else
         query_context = session.makeQueryContext();
     query_context->setCurrentQueryId(query_id);
     if (send_progress)

tests/queries/0_stateless/02303_query_kind.reference

@@ -0,0 +1,44 @@
+clickhouse-client --query_kind secondary_query -q explain plan header=1 select toString(dummy) as dummy from system.one group by dummy
+Expression ((Projection + Before ORDER BY))
+Header: dummy String
+  Aggregating
+  Header: toString(dummy) String
+    Expression (Before GROUP BY)
+    Header: toString(dummy) String
+      SettingQuotaAndLimits (Set limits and quota after reading from storage)
+      Header: dummy UInt8
+        ReadFromStorage (SystemOne)
+        Header: dummy UInt8
+clickhouse-local --query_kind secondary_query -q explain plan header=1 select toString(dummy) as dummy from system.one group by dummy
+Expression ((Projection + Before ORDER BY))
+Header: dummy String
+  Aggregating
+  Header: toString(dummy) String
+    Expression (Before GROUP BY)
+    Header: toString(dummy) String
+      SettingQuotaAndLimits (Set limits and quota after reading from storage)
+      Header: dummy UInt8
+        ReadFromStorage (SystemOne)
+        Header: dummy UInt8
+clickhouse-client --query_kind initial_query -q explain plan header=1 select toString(dummy) as dummy from system.one group by dummy
+Expression ((Projection + Before ORDER BY))
+Header: dummy String
+  Aggregating
+  Header: dummy UInt8
+    Expression (Before GROUP BY)
+    Header: dummy UInt8
+      SettingQuotaAndLimits (Set limits and quota after reading from storage)
+      Header: dummy UInt8
+        ReadFromStorage (SystemOne)
+        Header: dummy UInt8
+clickhouse-local --query_kind initial_query -q explain plan header=1 select toString(dummy) as dummy from system.one group by dummy
+Expression ((Projection + Before ORDER BY))
+Header: dummy String
+  Aggregating
+  Header: dummy UInt8
+    Expression (Before GROUP BY)
+    Header: dummy UInt8
+      SettingQuotaAndLimits (Set limits and quota after reading from storage)
+      Header: dummy UInt8
+        ReadFromStorage (SystemOne)
+        Header: dummy UInt8

tests/queries/0_stateless/02303_query_kind.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+CUR_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CUR_DIR"/../shell_config.sh
+
+function run_query()
+{
+    echo "clickhouse-client $*"
+    $CLICKHOUSE_CLIENT "$@"
+
+    echo "clickhouse-local $*"
+    $CLICKHOUSE_LOCAL "$@"
+}
+run_query --query_kind secondary_query -q "explain plan header=1 select toString(dummy) as dummy from system.one group by dummy"
+run_query --query_kind initial_query -q "explain plan header=1 select toString(dummy) as dummy from system.one group by dummy"