refactor clickhouse benchmark

2 files changed, 44 insertions, 24 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/click_bench.cpp b/ydb/public/lib/ydb_cli/commands/click_bench.cpp
index 9d76e89b7c..c65c5513e9 100644
--- a/ydb/public/lib/ydb_cli/commands/click_bench.cpp
+++ b/ydb/public/lib/ydb_cli/commands/click_bench.cpp
@@ -1,6 +1,7 @@
 #include <util/string/split.h>
 #include <util/stream/file.h>
 #include <util/string/strip.h>
+#include <util/string/join.h>
 #include <util/string/printf.h>
 #include <util/folder/pathsplit.h>
 #include <library/cpp/json/json_writer.h>
@@ -88,6 +89,39 @@ TClickHouseBench::TTestInfo TClickHouseBench::AnalyzeTestRuns(const TVector<TDur
     return info;
 }
 
+TString TBenchContext::PatchQuery(const TStringBuf& original) const {
+    TString result(original.data(), original.size());
+    if (EnablePushdown) {
+        result = "PRAGMA ydb.KqpPushOlapProcess = \"true\";\n" + result;
+    }
+    if (DisableLlvm) {
+        result = "PRAGMA ydb.EnableLlvm=\"false\";\n" + result;
+    }
+
+    std::vector<TStringBuf> lines;
+    for(auto& line : StringSplitter(result).Split('\n').SkipEmpty()) {
+        if (line.StartsWith("--")) {
+            continue;
+        }
+
+        lines.push_back(line);
+    }
+
+    return JoinSeq('\n', lines);
+}
+
+
+bool TBenchContext::NeedRun(const ui32 queryIdx) const {
+    if (QueriesToRun.size() && !QueriesToRun.contains(queryIdx)) {
+        return false;
+    }
+    if (QueriesToSkip.contains(queryIdx)) {
+        return false;
+    }
+    return true;
+}
+
+
 TClickBenchCommandInit::TClickBenchCommandInit()
     : TYdbCommand("init", {"i"}, "Initialize table")
 {}
@@ -355,3 +389,10 @@ std::pair<TString, TString> TStreamQueryRunner::ResultToYson(NTable::TScanQueryP
     writer.OnEndList();
     return {out.Str(), err_out.Str()};
 }
+
+TCommandClickBench::TCommandClickBench()
+    : TClientCommandTree("click_bench")
+{
+    AddCommand(std::make_unique<TClickBenchCommandRun>());
+    AddCommand(std::make_unique<TClickBenchCommandInit>());
+}
diff --git a/ydb/public/lib/ydb_cli/commands/click_bench.h b/ydb/public/lib/ydb_cli/commands/click_bench.h
index 69b2821889..ed472434cc 100644
--- a/ydb/public/lib/ydb_cli/commands/click_bench.h
+++ b/ydb/public/lib/ydb_cli/commands/click_bench.h
@@ -37,26 +37,8 @@ public:
         return externalQueries;
     }
 
-    TString PatchQuery(const TStringBuf& original) const {
-        TString result(original.data(), original.size());
-        if (EnablePushdown) {
-            result = "PRAGMA ydb.KqpPushOlapProcess = \"true\";\n" + result;
-        }
-        if (DisableLlvm) {
-            result = "PRAGMA ydb.EnableLlvm=\"false\";\n" + result;
-        }
-        return result;
-    }
-
-    bool NeedRun(const ui32 queryIdx) const {
-        if (QueriesToRun.size() && !QueriesToRun.contains(queryIdx)) {
-            return false;
-        }
-        if (QueriesToSkip.contains(queryIdx)) {
-            return false;
-        }
-        return true;
-    }
+    TString PatchQuery(const TStringBuf& original) const;
+    bool NeedRun(const ui32 queryIdx) const;
 };
 
 class TClickHouseBench {
@@ -129,9 +111,6 @@ private:
 
 class TCommandClickBench : public NYdb::NConsoleClient::TClientCommandTree {
 public:
-    TCommandClickBench() : TClientCommandTree("click_bench") {
-        AddCommand(std::make_unique<TClickBenchCommandRun>());
-        AddCommand(std::make_unique<TClickBenchCommandInit>());
-    }
+    TCommandClickBench();
 };