Merge pull request #9416 from CurtizJ/benchmark-query-id

Allow to set query_id in clickhouse-benchmark

dbms/programs/benchmark/Benchmark.cpp

@@ -59,11 +59,11 @@ public:
             bool cumulative_, bool secure_, const String & default_database_,
             const String & user_, const String & password_, const String & stage,
             bool randomize_, size_t max_iterations_, double max_time_,
-            const String & json_path_, size_t confidence_, const Settings & settings_)
+            const String & json_path_, size_t confidence_, const String & query_id_, const Settings & settings_)
         :
         concurrency(concurrency_), delay(delay_), queue(concurrency), randomize(randomize_),
         cumulative(cumulative_), max_iterations(max_iterations_), max_time(max_time_),
-        confidence(confidence_), json_path(json_path_), settings(settings_),
+        json_path(json_path_), confidence(confidence_), query_id(query_id_), settings(settings_),
         global_context(Context::createGlobal()), pool(concurrency)
     {
         const auto secure = secure_ ? Protocol::Secure::Enable : Protocol::Secure::Disable;
@@ -145,8 +145,9 @@ private:
     bool cumulative;
     size_t max_iterations;
     double max_time;
-    size_t confidence;
     String json_path;
+    size_t confidence;
+    std::string query_id;
     Settings settings;
     Context global_context;
     QueryProcessingStage::Enum query_processing_stage;
@@ -367,6 +368,8 @@ private:
         RemoteBlockInputStream stream(
             *(*connection_entries[connection_index]),
             query, {}, global_context, &settings, nullptr, Scalars(), Tables(), query_processing_stage);
+        if (!query_id.empty())
+            stream.setQueryId(query_id);
 
         Progress progress;
         stream.setProgressCallback([&progress](const Progress & value) { progress.incrementPiecewiseAtomically(value); });
@@ -535,6 +538,7 @@ int mainEntryClickHouseBenchmark(int argc, char ** argv)
             ("database",      value<std::string>()->default_value("default"),   "")
             ("stacktrace",                                                      "print stack traces of exceptions")
             ("confidence",    value<size_t>()->default_value(5),                "set the level of confidence for T-test [0=80%, 1=90%, 2=95%, 3=98%, 4=99%, 5=99.5%(default)")
+            ("query_id",      value<std::string>()->default_value(""),         "")
         ;
 
         Settings settings;
@@ -573,6 +577,7 @@ int mainEntryClickHouseBenchmark(int argc, char ** argv)
             options["timelimit"].as<double>(),
             options["json"].as<std::string>(),
             options["confidence"].as<size_t>(),
+            options["query_id"].as<std::string>(),
             settings);
         return benchmark.run();
     }

dbms/tests/queries/0_stateless/01088_benchmark_query_id.reference

+20

dbms/tests/queries/0_stateless/01088_benchmark_query_id.sh

+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. $CURDIR/../shell_config.sh
+
+QUERY_ID=$RANDOM
+$CLICKHOUSE_BENCHMARK <<< "SELECT 1" --query_id $QUERY_ID -i 10 2>/dev/null
+$CLICKHOUSE_CLIENT -q "SYSTEM FLUSH LOGS"
+$CLICKHOUSE_CLIENT -q "SELECT count() FROM system.query_log WHERE query_id='$QUERY_ID'"

dbms/tests/queries/shell_config.sh

@@ -15,6 +15,7 @@ export CLICKHOUSE_CLIENT=${CLICKHOUSE_CLIENT:="$CLICKHOUSE_CLIENT_BINARY ${CLICK
 [ -x "${CLICKHOUSE_BINARY}" ] && CLICKHOUSE_LOCAL=${CLICKHOUSE_LOCAL:="${CLICKHOUSE_BINARY} local"}
 export CLICKHOUSE_LOCAL=${CLICKHOUSE_LOCAL:="${CLICKHOUSE_BINARY}-local"}
 export CLICKHOUSE_OBFUSCATOR=${CLICKHOUSE_OBFUSCATOR:="${CLICKHOUSE_BINARY}-obfuscator"}
+export CLICKHOUSE_BENCHMARK=${CLICKHOUSE_BENCHMARK:="${CLICKHOUSE_BINARY}-benchmark"}
 
 export CLICKHOUSE_CONFIG=${CLICKHOUSE_CONFIG:="/etc/clickhouse-server/config.xml"}
 export CLICKHOUSE_CONFIG_CLIENT=${CLICKHOUSE_CONFIG_CLIENT:="/etc/clickhouse-client/config.xml"}