Query workload over query service (#9564)

2 files changed, 11 insertions, 9 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/query_workload.cpp b/ydb/public/lib/ydb_cli/commands/query_workload.cpp
index 09446cda4ff..9aadd3149a1 100644
--- a/ydb/public/lib/ydb_cli/commands/query_workload.cpp
+++ b/ydb/public/lib/ydb_cli/commands/query_workload.cpp
@@ -1,6 +1,7 @@
 #include "query_workload.h"
 
 #include <ydb/public/lib/ydb_cli/commands/ydb_common.h>
+#include <ydb/public/sdk/cpp/client/ydb_query/client.h>
 #include <library/cpp/histogram/hdr/histogram.h>
 #include <util/system/thread.h>
 #include <util/system/mutex.h>
@@ -77,11 +78,11 @@ TCommandQueryWorkload::TCommandQueryWorkload()
 }
 
 TCommandQueryWorkloadRun::TCommandQueryWorkloadRun()
-    : TYdbOperationCommand("run", {}, "Run YDB query workload")
+    : TYdbSimpleCommand("run", {}, "Run YDB query workload")
 {}
 
 void TCommandQueryWorkloadRun::Config(TConfig& config) {
-    TYdbOperationCommand::Config(config);
+    TYdbSimpleCommand::Config(config);
     config.Opts->AddLongOption('q', "query", "Query to execute") .RequiredArgument("[String]").StoreResult(&Query);
     config.Opts->AddLongOption('t', "threads", "Number of parallel threads; 1 if not specified").DefaultValue(1).StoreResult(&Threads);
     config.Opts->AddLongOption('d', "delay", "Interval delay in seconds; 1 if not specified").DefaultValue(1).StoreResult(&IntervalSeconds);
@@ -102,14 +103,15 @@ int TCommandQueryWorkloadRun::Run(TConfig& config) {
         pool.SafeAddFunc([&] {
             try {
                 TDriver driver = CreateDriver(config);
-                NScripting::TScriptingClient client(driver);
+                NQuery::TQueryClient client(driver);
 
-                NScripting::TExecuteYqlRequestSettings settings;
-                settings.CollectQueryStats(NTable::ECollectQueryStatsMode::Basic);
+                NQuery::TExecuteQuerySettings settings;
+                settings.StatsMode(NQuery::EStatsMode::Basic);
 
                 while (!IsInterrupted() && !ThreadTerminated.load()) {
-                    auto asyncResult = client.StreamExecuteYqlScript(
+                    auto asyncResult = client.StreamExecuteQuery(
                         Query,
+                        NQuery::TTxControl::NoTx(),
                         FillSettings(settings)
                     );
 
@@ -129,8 +131,8 @@ int TCommandQueryWorkloadRun::Run(TConfig& config) {
                             ThrowOnError(streamPart);
                         }
 
-                        if (streamPart.HasQueryStats()) {
-                            const auto& queryStats = streamPart.GetQueryStats();
+                        if (streamPart.GetStats()) {
+                            const auto& queryStats = streamPart.GetStats().GetRef();
                             local_duration += queryStats.GetTotalDuration();
                         }
                     }
diff --git a/ydb/public/lib/ydb_cli/commands/query_workload.h b/ydb/public/lib/ydb_cli/commands/query_workload.h
index 033bb0e8844..0c03b9cd3a3 100644
--- a/ydb/public/lib/ydb_cli/commands/query_workload.h
+++ b/ydb/public/lib/ydb_cli/commands/query_workload.h
@@ -11,7 +11,7 @@ public:
     TCommandQueryWorkload();
 };
 
-class TCommandQueryWorkloadRun : public TYdbOperationCommand, public TInterruptibleCommand
+class TCommandQueryWorkloadRun : public TYdbSimpleCommand, public TInterruptibleCommand
 {
 public:
     TCommandQueryWorkloadRun();