From 960440e59aec63c6ebc9f5e590d19bfd8c333a63 Mon Sep 17 00:00:00 2001
From: Anton Popov <anton@clickhouse.com>
Date: Fri, 3 May 2024 16:20:03 +0000
Subject: [PATCH] add option query_id_prefix to clickhouse-benchmark

---
 programs/benchmark/Benchmark.cpp                      | 10 ++++++++++
 .../03143_benchmark_query_id_prefix.reference         |  1 +
 .../0_stateless/03143_benchmark_query_id_prefix.sh    | 11 +++++++++++
 3 files changed, 22 insertions(+)
 create mode 100644 tests/queries/0_stateless/03143_benchmark_query_id_prefix.reference
 create mode 100755 tests/queries/0_stateless/03143_benchmark_query_id_prefix.sh

diff --git a/programs/benchmark/Benchmark.cpp b/programs/benchmark/Benchmark.cpp
index eecc352d073..48dca82eb2b 100644
--- a/programs/benchmark/Benchmark.cpp
+++ b/programs/benchmark/Benchmark.cpp
@@ -7,6 +7,8 @@
 #include <random>
 #include <string_view>
 #include <pcg_random.hpp>
+#include <Poco/UUID.h>
+#include <Poco/UUIDGenerator.h>
 #include <Poco/Util/Application.h>
 #include <Common/Stopwatch.h>
 #include <Common/ThreadPool.h>
@@ -80,6 +82,7 @@ public:
             double max_time_,
             size_t confidence_,
             const String & query_id_,
+            const String & query_id_prefix_,
             const String & query_to_execute_,
             size_t max_consecutive_errors_,
             bool continue_on_errors_,
@@ -98,6 +101,7 @@ public:
         max_time(max_time_),
         confidence(confidence_),
         query_id(query_id_),
+        query_id_prefix(query_id_prefix_),
         query_to_execute(query_to_execute_),
         continue_on_errors(continue_on_errors_),
         max_consecutive_errors(max_consecutive_errors_),
@@ -205,6 +209,7 @@ private:
     double max_time;
     size_t confidence;
     String query_id;
+    String query_id_prefix;
     String query_to_execute;
     bool continue_on_errors;
     size_t max_consecutive_errors;
@@ -463,8 +468,11 @@ private:
 
         RemoteQueryExecutor executor(
             *entry, query, {}, global_context, nullptr, Scalars(), Tables(), query_processing_stage);
+
         if (!query_id.empty())
             executor.setQueryId(query_id);
+        else if (!query_id_prefix.empty())
+            executor.setQueryId(query_id_prefix + "_" + Poco::UUIDGenerator().createRandom().toString());
 
         Progress progress;
         executor.setProgressCallback([&progress](const Progress & value) { progress.incrementPiecewiseAtomically(value); });
@@ -617,6 +625,7 @@ int mainEntryClickHouseBenchmark(int argc, char ** argv)
             ("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(""), "")
+            ("query_id_prefix", value<std::string>()->default_value(""), "")
             ("max-consecutive-errors", value<size_t>()->default_value(0), "set number of allowed consecutive errors")
             ("ignore-error,continue_on_errors", "continue testing even if a query fails")
             ("reconnect", "establish new connection for every query")
@@ -671,6 +680,7 @@ int mainEntryClickHouseBenchmark(int argc, char ** argv)
             options["timelimit"].as<double>(),
             options["confidence"].as<size_t>(),
             options["query_id"].as<std::string>(),
+            options["query_id_prefix"].as<std::string>(),
             options["query"].as<std::string>(),
             options["max-consecutive-errors"].as<size_t>(),
             options.count("ignore-error"),
diff --git a/tests/queries/0_stateless/03143_benchmark_query_id_prefix.reference b/tests/queries/0_stateless/03143_benchmark_query_id_prefix.reference
new file mode 100644
index 00000000000..d7a23e16ea5
--- /dev/null
+++ b/tests/queries/0_stateless/03143_benchmark_query_id_prefix.reference
@@ -0,0 +1 @@
+0	100
diff --git a/tests/queries/0_stateless/03143_benchmark_query_id_prefix.sh b/tests/queries/0_stateless/03143_benchmark_query_id_prefix.sh
new file mode 100755
index 00000000000..ed68b443c9a
--- /dev/null
+++ b/tests/queries/0_stateless/03143_benchmark_query_id_prefix.sh
@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CURDIR"/../shell_config.sh
+
+query_id_prefix=${CLICKHOUSE_DATABASE}_test_benchmark
+$CLICKHOUSE_BENCHMARK -i 100 -c 8 <<< "SELECT 1" --query_id_prefix $query_id_prefix 2>/dev/null
+
+$CLICKHOUSE_CLIENT --query "SYSTEM FLUSH LOGS"
+$CLICKHOUSE_CLIENT --query "SELECT countIf(query_id = '$query_id_prefix'), countIf(query_id LIKE '$query_id_prefix%') FROM system.query_log WHERE current_database = currentDatabase() AND type = 'QueryFinish'"