Support progress stats for workload (#13725)

Author: shnikd

include/ydb-cpp-sdk/client/query/client.h

@@ -210,7 +210,10 @@ class TExecuteQueryPart : public TStreamPartStatus {
     const TResultSet& GetResultSet() const { return *ResultSet_; }
     TResultSet ExtractResultSet() { return std::move(*ResultSet_); }
 
+    bool HasStats() const { return Stats_.has_value(); }
     const std::optional<TExecStats>& GetStats() const { return Stats_; }
+    TExecStats ExtractStats() const { return std::move(*Stats_); }
+    
     const std::optional<TTransaction>& GetTransaction() const { return Transaction_; }
 
     TExecuteQueryPart(TStatus&& status, std::optional<TExecStats>&& queryStats, std::optional<TTransaction>&& tx)

include/ydb-cpp-sdk/client/query/query.h

@@ -76,6 +76,7 @@ struct TExecuteQuerySettings : public TRequestSettings<TExecuteQuerySettings> {
     FLUENT_SETTING_DEFAULT(EStatsMode, StatsMode, EStatsMode::None);
     FLUENT_SETTING_OPTIONAL(bool, ConcurrentResultSets);
     FLUENT_SETTING(std::string, ResourcePool);
+    FLUENT_SETTING_OPTIONAL(std::chrono::milliseconds, StatsCollectPeriod);
 };
 
 struct TBeginTxSettings : public TRequestSettings<TBeginTxSettings> {};

src/api/protos/ydb_query.proto

@@ -174,7 +174,13 @@ message ExecuteQueryRequest {
     // Allows to set size limitation (in bytes) for one result part
     int64 response_part_limit_bytes = 9 [(Ydb.value) = "[0; 33554432]"];
 
-    string pool_id = 10; // Workload manager pool id
+    // Workload manager pool id
+    string pool_id = 10;
+
+    // Time interval for sending periodical query statistics.
+    // When query statistics are enabled (stats_mode != STATS_MODE_NONE), by default statistics will be sent only once after query execution is finished.
+    // In case when stats_period_ms is specified and is non-zero, query statistics will be additionally sent every stats_period_ms milliseconds beginning from the start of query execution.
+    int64 stats_period_ms = 11 [(Ydb.value) = ">= 0"];
 }
 
 message ResultSetMeta {

src/client/query/impl/exec_query.cpp

@@ -246,6 +246,10 @@ TFuture<std::pair<TPlainStatus, TExecuteQueryProcessorPtr>> StreamExecuteQueryIm
         request.set_response_part_limit_bytes(*settings.OutputChunkMaxSize_);
     }
 
+    if (settings.StatsCollectPeriod_) {
+        request.set_stats_period_ms((*settings.StatsCollectPeriod_).count());
+    }
+
     if (txControl.HasTx()) {
         auto requestTxControl = request.mutable_tx_control();
         requestTxControl->set_commit_tx(txControl.CommitTx_);