Merge pull request #16843 from azat/clickhouse-client-stage

Add --stage option for clickhouse-client

programs/benchmark/Benchmark.cpp

@@ -49,7 +49,6 @@ using Ports = std::vector<UInt16>;
 namespace ErrorCodes
 {
     extern const int CANNOT_BLOCK_SIGNAL;
-    extern const int BAD_ARGUMENTS;
     extern const int EMPTY_DATA_PASSED;
 }
 
@@ -103,17 +102,7 @@ public:
         /// (example: when using stage = 'with_mergeable_state')
         registerAggregateFunctions();
 
-        if (stage == "complete")
-            query_processing_stage = QueryProcessingStage::Complete;
-        else if (stage == "fetch_columns")
-            query_processing_stage = QueryProcessingStage::FetchColumns;
-        else if (stage == "with_mergeable_state")
-            query_processing_stage = QueryProcessingStage::WithMergeableState;
-        else if (stage == "with_mergeable_state_after_aggregation")
-            query_processing_stage = QueryProcessingStage::WithMergeableStateAfterAggregation;
-        else
-            throw Exception("Unknown query processing stage: " + stage, ErrorCodes::BAD_ARGUMENTS);
-
+        query_processing_stage = QueryProcessingStage::fromString(stage);
     }
 
     void initialize(Poco::Util::Application & self [[maybe_unused]]) override

programs/client/Client.cpp

@@ -224,6 +224,7 @@ private:
 
     /// We will format query_id in interactive mode in various ways, the default is just to print Query id: ...
     std::vector<std::pair<String, String>> query_id_formats;
+    QueryProcessingStage::Enum query_processing_stage;
 
     void initialize(Poco::Util::Application & self) override
     {
@@ -1444,7 +1445,7 @@ private:
                     connection_parameters.timeouts,
                     query_to_send,
                     context.getCurrentQueryId(),
-                    QueryProcessingStage::Complete,
+                    query_processing_stage,
                     &context.getSettingsRef(),
                     &context.getClientInfo(),
                     true);
@@ -1485,7 +1486,7 @@ private:
             connection_parameters.timeouts,
             query_to_send,
             context.getCurrentQueryId(),
-            QueryProcessingStage::Complete,
+            query_processing_stage,
             &context.getSettingsRef(),
             &context.getClientInfo(),
             true);
@@ -2309,6 +2310,7 @@ public:
             ("password", po::value<std::string>()->implicit_value("\n", ""), "password")
             ("ask-password", "ask-password")
             ("quota_key", po::value<std::string>(), "A string to differentiate quotas when the user have keyed quotas configured on server")
+            ("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")
             ("query_id", po::value<std::string>(), "query_id")
             ("query,q", po::value<std::string>(), "query")
             ("database,d", po::value<std::string>(), "database")
@@ -2433,6 +2435,8 @@ public:
         if (options.count("config-file") && options.count("config"))
             throw Exception("Two or more configuration files referenced in arguments", ErrorCodes::BAD_ARGUMENTS);
 
+        query_processing_stage = QueryProcessingStage::fromString(options["stage"].as<std::string>());
+
         /// Save received data into the internal config.
         if (options.count("config-file"))
             config().setString("config-file", options["config-file"].as<std::string>());

src/Core/QueryProcessingStage.cpp

@@ -0,0 +1,35 @@
+#include <Core/QueryProcessingStage.h>
+#include <Common/Exception.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int BAD_ARGUMENTS;
+}
+
+namespace QueryProcessingStage
+{
+
+    Enum fromString(const std::string & stage_string)
+    {
+        Enum stage;
+
+        if (stage_string == "complete")
+            stage = Complete;
+        else if (stage_string == "fetch_columns")
+            stage = FetchColumns;
+        else if (stage_string == "with_mergeable_state")
+            stage = WithMergeableState;
+        else if (stage_string == "with_mergeable_state_after_aggregation")
+            stage = WithMergeableStateAfterAggregation;
+        else
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unknown query processing stage: {}", stage_string);
+
+        return stage;
+    }
+
+}
+
+}

src/Core/QueryProcessingStage.h

@@ -43,6 +43,14 @@ namespace QueryProcessingStage
             ? data[stage]
             : "Unknown stage";
     }
+
+    /// This methid is used for the program options,
+    /// hence it accept under_score notation for stage:
+    /// - complete
+    /// - fetch_columns
+    /// - with_mergeable_state
+    /// - with_mergeable_state_after_aggregation
+    Enum fromString(const std::string & stage_string);
 }
 
 }

src/Core/ya.make

@@ -30,6 +30,7 @@ SRCS(
     MySQL/PacketsReplication.cpp
     NamesAndTypes.cpp
     PostgreSQLProtocol.cpp
+    QueryProcessingStage.cpp
     Settings.cpp
     SettingsEnums.cpp
     SettingsFields.cpp

tests/queries/0_stateless/01561_clickhouse_client_stage.reference

@@ -0,0 +1,15 @@
+execute: default
+"foo"
+1
+execute: --stage fetch_columns
+"dummy"
+0
+execute: --stage with_mergeable_state
+"1"
+1
+execute: --stage with_mergeable_state_after_aggregation
+"1"
+1
+execute: --stage complete
+"foo"
+1

tests/queries/0_stateless/01561_clickhouse_client_stage.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../shell_config.sh
+
+function execute_query()
+{
+    if [ $# -eq 0 ]; then
+        echo "execute: default"
+    else
+        echo "execute: $*"
+    fi
+    ${CLICKHOUSE_CLIENT} "$@" --format CSVWithNames -q "SELECT 1 AS foo"
+}
+
+execute_query # default -- complete
+execute_query --stage fetch_columns
+execute_query --stage with_mergeable_state
+execute_query --stage with_mergeable_state_after_aggregation
+execute_query --stage complete