Fix Abort on dry run benchmarks (#12886)

Author: iddqdex

ydb/apps/ydb/CHANGELOG.md

@@ -1,3 +1,4 @@
+* Fixed a bug where `ydb workload * run` command could crash in `--dry-run` mode.
 * Added support for views in local backups: `ydb tools dump` and `ydb tools restore`. Views are backed up as `CREATE VIEW` queries saved in the `create_view.sql` files, which can be executed to recreate the original views.
 * Replaced option `--query-settings` by `--query-prefix` one in `ydb workload <workload> run`.
 * Added new options to `ydb workload topic`: --tx-commit-interval and --tx-commit-messages, allowing you to specify commit interval either in milliseconds or in number of messages written. 

ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp

@@ -286,7 +286,7 @@ void CollectStats(TPrettyTable& table, IOutputStream* csv, NJson::TJsonValue* js
 }
 
 template <typename TClient>
-bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen) {
+bool TWorkloadCommandBenchmark::RunBench(TClient* client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen) {
     using namespace BenchmarkUtils;
     TOFStream outFStream{OutFilePath};
     TPrettyTable statTable(ColumnNames);
@@ -344,7 +344,11 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
         if (PlanFileName) {
             TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined", "undefined", "undefined");
             try {
-                res = Explain(query, client, GetDeadline());
+                if (client) {
+                    res = Explain(query, *client, GetDeadline());
+                } else {
+                    res = TQueryBenchmarkResult::Result(TQueryBenchmarkResult::TRawResults(), TDuration::Zero(), "", "");
+                }
             } catch (...) {
                 res = TQueryBenchmarkResult::Error(CurrentExceptionMessage(), "", "");
             }
@@ -355,7 +359,11 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
             auto t1 = TInstant::Now();
             TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined", "undefined", "undefined");
             try {
-                res = Execute(query, client, GetDeadline());
+                if (client) {
+                    res = Execute(query, *client, GetDeadline());
+                } else {
+                    res = TQueryBenchmarkResult::Result(TQueryBenchmarkResult::TRawResults(), TDuration::Zero(), "", "");
+                }
             } catch (...) {
                 res = TQueryBenchmarkResult::Error(CurrentExceptionMessage(), "", "");
             }
@@ -526,10 +534,10 @@ BenchmarkUtils::TQueryBenchmarkDeadline TWorkloadCommandBenchmark::GetDeadline()
 
 int TWorkloadCommandBenchmark::DoRun(NYdbWorkload::IWorkloadQueryGenerator& workloadGen, TConfig& /*config*/) {
     if (QueryExecuterType == "scan") {
-        return !RunBench(*TableClient, workloadGen);
+        return !RunBench(TableClient.Get(), workloadGen);
     }
     if (QueryExecuterType == "generic") {
-        return !RunBench(*QueryClient, workloadGen);
+        return !RunBench(QueryClient.Get(), workloadGen);
     }
     ythrow yexception() << "Incorrect executer type. Available options: \"scan\", \"generic\"." << Endl;
 }

ydb/public/lib/ydb_cli/commands/ydb_benchmark.h

@@ -20,7 +20,7 @@ class TWorkloadCommandBenchmark final: public TWorkloadCommandBase {
     bool NeedRun(const ui32 queryIdx) const;
 
     template <typename TClient>
-    bool RunBench(TClient& client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen);
+    bool RunBench(TClient* client, NYdbWorkload::IWorkloadQueryGenerator& workloadGen);
     void SavePlans(const BenchmarkUtils::TQueryBenchmarkResult& res, ui32 queryNum, const TStringBuf name) const;
     void PrintResult(const BenchmarkUtils::TQueryBenchmarkResult& res, IOutputStream& out, const std::string& expected) const;
     BenchmarkUtils::TQueryBenchmarkDeadline GetDeadline() const;