add api to interact with query stats (#44507)

Co-authored-by: Shfdis <[email protected]>

4 files changed, 313 insertions, 2 deletions

diff --git a/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/stats.h b/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/stats.h
index 57754e90d91..64b2b7caf78 100644
--- a/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/stats.h
+++ b/ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/stats.h
@@ -5,9 +5,14 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <vector>
 
 namespace Ydb::TableStats {
+    class CompilationStats;
+    class OperationStats;
+    class QueryPhaseStats;
     class QueryStats;
+    class TableAccessStats;
 }
 
 namespace NYdb::inline Dev {
@@ -16,6 +21,72 @@ namespace NYdb::inline Dev {
 
 namespace NYdb::inline Dev::NQuery {
 
+class TOperationStats {
+public:
+    explicit TOperationStats(const Ydb::TableStats::OperationStats& proto);
+
+    uint64_t GetRows() const;
+    uint64_t GetBytes() const;
+
+private:
+    uint64_t Rows_ = 0;
+    uint64_t Bytes_ = 0;
+};
+
+class TTableAccessStats {
+public:
+    explicit TTableAccessStats(const Ydb::TableStats::TableAccessStats& proto);
+
+    const std::string& GetName() const;
+    const TOperationStats& GetReads() const;
+    const TOperationStats& GetUpdates() const;
+    const TOperationStats& GetDeletes() const;
+    uint64_t GetPartitionsCount() const;
+
+private:
+    std::string Name_;
+    TOperationStats Reads_;
+    TOperationStats Updates_;
+    TOperationStats Deletes_;
+    uint64_t PartitionsCount_ = 0;
+};
+
+class TQueryPhaseStats {
+public:
+    explicit TQueryPhaseStats(const Ydb::TableStats::QueryPhaseStats& proto);
+
+    uint64_t GetDurationUs() const;
+    TDuration GetDuration() const;
+    uint64_t GetCpuTimeUs() const;
+    TDuration GetCpuTime() const;
+    uint64_t GetAffectedShards() const;
+    bool IsLiteralPhase() const;
+    const std::vector<TTableAccessStats>& GetTableAccess() const;
+
+private:
+    uint64_t DurationUs_ = 0;
+    uint64_t CpuTimeUs_ = 0;
+    uint64_t AffectedShards_ = 0;
+    bool LiteralPhase_ = false;
+    std::vector<TTableAccessStats> TableAccess_;
+};
+
+class TCompilationStats {
+public:
+    explicit TCompilationStats(const Ydb::TableStats::CompilationStats& proto);
+
+    bool IsFromCache() const;
+    uint64_t GetDurationUs() const;
+    TDuration GetDuration() const;
+    uint64_t GetCpuTimeUs() const;
+    TDuration GetCpuTime() const;
+
+private:
+    bool FromCache_ = false;
+    uint64_t DurationUs_ = 0;
+    uint64_t CpuTimeUs_ = 0;
+};
+
 class TExecStats {
     friend class NYdb::TProtoAccessor;
 
@@ -27,13 +98,20 @@ public:
 
     std::string ToString(bool withPlan = false) const;
 
+    uint64_t GetProcessCpuTimeUs() const;
+    uint64_t GetTotalDurationUs() const;
+    uint64_t GetTotalCpuTimeUs() const;
     std::optional<std::string> GetPlan() const;
     std::optional<std::string> GetAst() const;
     std::optional<std::string> GetMeta() const;
 
+    TDuration GetProcessCpuTime() const;
     TDuration GetTotalDuration() const;
     TDuration GetTotalCpuTime() const;
 
+    std::vector<TQueryPhaseStats> GetQueryPhases() const;
+    std::optional<TCompilationStats> GetCompilation() const;
+
 private:
     const Ydb::TableStats::QueryStats& GetProto() const;
 
diff --git a/ydb/public/sdk/cpp/src/client/query/stats.cpp b/ydb/public/sdk/cpp/src/client/query/stats.cpp
index 0e0a6f161a5..76913abde75 100644
--- a/ydb/public/sdk/cpp/src/client/query/stats.cpp
+++ b/ydb/public/sdk/cpp/src/client/query/stats.cpp
@@ -15,6 +15,113 @@ public:
     Ydb::TableStats::QueryStats Proto;
 };
 
+TOperationStats::TOperationStats(const Ydb::TableStats::OperationStats& proto)
+    : Rows_(proto.rows())
+    , Bytes_(proto.bytes())
+{}
+
+uint64_t TOperationStats::GetRows() const {
+    return Rows_;
+}
+
+uint64_t TOperationStats::GetBytes() const {
+    return Bytes_;
+}
+
+TTableAccessStats::TTableAccessStats(const Ydb::TableStats::TableAccessStats& proto)
+    : Name_(proto.name())
+    , Reads_(proto.reads())
+    , Updates_(proto.updates())
+    , Deletes_(proto.deletes())
+    , PartitionsCount_(proto.partitions_count())
+{}
+
+const std::string& TTableAccessStats::GetName() const {
+    return Name_;
+}
+
+const TOperationStats& TTableAccessStats::GetReads() const {
+    return Reads_;
+}
+
+const TOperationStats& TTableAccessStats::GetUpdates() const {
+    return Updates_;
+}
+
+const TOperationStats& TTableAccessStats::GetDeletes() const {
+    return Deletes_;
+}
+
+uint64_t TTableAccessStats::GetPartitionsCount() const {
+    return PartitionsCount_;
+}
+
+TQueryPhaseStats::TQueryPhaseStats(const Ydb::TableStats::QueryPhaseStats& proto)
+    : DurationUs_(proto.duration_us())
+    , CpuTimeUs_(proto.cpu_time_us())
+    , AffectedShards_(proto.affected_shards())
+    , LiteralPhase_(proto.literal_phase())
+{
+    TableAccess_.reserve(proto.table_access().size());
+    for (const auto& tableAccess : proto.table_access()) {
+        TableAccess_.emplace_back(tableAccess);
+    }
+}
+
+uint64_t TQueryPhaseStats::GetDurationUs() const {
+    return DurationUs_;
+}
+
+TDuration TQueryPhaseStats::GetDuration() const {
+    return TDuration::MicroSeconds(DurationUs_);
+}
+
+uint64_t TQueryPhaseStats::GetCpuTimeUs() const {
+    return CpuTimeUs_;
+}
+
+TDuration TQueryPhaseStats::GetCpuTime() const {
+    return TDuration::MicroSeconds(CpuTimeUs_);
+}
+
+uint64_t TQueryPhaseStats::GetAffectedShards() const {
+    return AffectedShards_;
+}
+
+bool TQueryPhaseStats::IsLiteralPhase() const {
+    return LiteralPhase_;
+}
+
+const std::vector<TTableAccessStats>& TQueryPhaseStats::GetTableAccess() const {
+    return TableAccess_;
+}
+
+TCompilationStats::TCompilationStats(const Ydb::TableStats::CompilationStats& proto)
+    : FromCache_(proto.from_cache())
+    , DurationUs_(proto.duration_us())
+    , CpuTimeUs_(proto.cpu_time_us())
+{}
+
+bool TCompilationStats::IsFromCache() const {
+    return FromCache_;
+}
+
+uint64_t TCompilationStats::GetDurationUs() const {
+    return DurationUs_;
+}
+
+TDuration TCompilationStats::GetDuration() const {
+    return TDuration::MicroSeconds(DurationUs_);
+}
+
+uint64_t TCompilationStats::GetCpuTimeUs() const {
+    return CpuTimeUs_;
+}
+
+TDuration TCompilationStats::GetCpuTime() const {
+    return TDuration::MicroSeconds(CpuTimeUs_);
+}
+
 TExecStats::TExecStats(const Ydb::TableStats::QueryStats& proto) {
     Impl_ = std::make_shared<TImpl>();
     Impl_->Proto = proto;
@@ -39,6 +146,18 @@ std::string TExecStats::ToString(bool withPlan) const {
     return res;
 }
 
+uint64_t TExecStats::GetProcessCpuTimeUs() const {
+    return Impl_->Proto.process_cpu_time_us();
+}
+
+uint64_t TExecStats::GetTotalDurationUs() const {
+    return Impl_->Proto.total_duration_us();
+}
+
+uint64_t TExecStats::GetTotalCpuTimeUs() const {
+    return Impl_->Proto.total_cpu_time_us();
+}
+
 std::optional<std::string> TExecStats::GetPlan() const {
     auto proto = Impl_->Proto;
 
@@ -69,6 +188,10 @@ std::optional<std::string> TExecStats::GetMeta() const {
     return proto.query_meta();
 }
 
+TDuration TExecStats::GetProcessCpuTime() const {
+    return TDuration::MicroSeconds(Impl_->Proto.process_cpu_time_us());
+}
+
 TDuration TExecStats::GetTotalDuration() const {
     return TDuration::MicroSeconds(Impl_->Proto.total_duration_us());
 }
@@ -77,6 +200,23 @@ TDuration TExecStats::GetTotalCpuTime() const {
     return TDuration::MicroSeconds(Impl_->Proto.total_cpu_time_us());
 }
 
+std::vector<TQueryPhaseStats> TExecStats::GetQueryPhases() const {
+    std::vector<TQueryPhaseStats> phases;
+    phases.reserve(Impl_->Proto.query_phases().size());
+    for (const auto& phase : Impl_->Proto.query_phases()) {
+        phases.emplace_back(phase);
+    }
+    return phases;
+}
+
+std::optional<TCompilationStats> TExecStats::GetCompilation() const {
+    const auto& proto = Impl_->Proto;
+    if (!proto.has_compilation()) {
+        return {};
+    }
+    return TCompilationStats(proto.compilation());
+}
+
 const Ydb::TableStats::QueryStats& TExecStats::GetProto() const {
     return Impl_->Proto;
 }
diff --git a/ydb/public/sdk/cpp/tests/unit/client/query/query_stats_ut.cpp b/ydb/public/sdk/cpp/tests/unit/client/query/query_stats_ut.cpp
new file mode 100644
index 00000000000..54bc83f7064
--- /dev/null
+++ b/ydb/public/sdk/cpp/tests/unit/client/query/query_stats_ut.cpp
@@ -0,0 +1,91 @@
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/stats.h>
+
+#include <ydb/public/api/protos/ydb_query_stats.pb.h>
+
+#include <library/cpp/testing/unittest/registar.h>
+
+using namespace NYdb::NQuery;
+
+Y_UNIT_TEST_SUITE(QueryStats) {
+    Y_UNIT_TEST(MapsProtoValues) {
+        Ydb::TableStats::QueryStats proto;
+        proto.set_process_cpu_time_us(10);
+        proto.set_total_duration_us(20);
+        proto.set_total_cpu_time_us(30);
+        proto.set_query_plan("plan");
+        proto.set_query_ast("ast");
+        proto.set_query_meta("meta");
+
+        auto* compilation = proto.mutable_compilation();
+        compilation->set_from_cache(true);
+        compilation->set_duration_us(40);
+        compilation->set_cpu_time_us(50);
+
+        auto* phase = proto.add_query_phases();
+        phase->set_duration_us(60);
+        phase->set_cpu_time_us(70);
+        phase->set_affected_shards(80);
+        phase->set_literal_phase(true);
+
+        auto* table = phase->add_table_access();
+        table->set_name("table");
+        table->set_partitions_count(90);
+        table->mutable_reads()->set_rows(100);
+        table->mutable_reads()->set_bytes(110);
+        table->mutable_updates()->set_rows(120);
+        table->mutable_updates()->set_bytes(130);
+        table->mutable_deletes()->set_rows(140);
+        table->mutable_deletes()->set_bytes(150);
+
+        const TExecStats stats(proto);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetProcessCpuTimeUs(), 10);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetProcessCpuTime().MicroSeconds(), 10);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalDurationUs(), 20);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalDuration().MicroSeconds(), 20);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalCpuTimeUs(), 30);
+        UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalCpuTime().MicroSeconds(), 30);
+        const auto plan = stats.GetPlan();
+        const auto ast = stats.GetAst();
+        const auto meta = stats.GetMeta();
+        UNIT_ASSERT_VALUES_EQUAL(*plan, "plan");
+        UNIT_ASSERT_VALUES_EQUAL(*ast, "ast");
+        UNIT_ASSERT_VALUES_EQUAL(*meta, "meta");
+
+        const auto compilationStats = stats.GetCompilation();
+        UNIT_ASSERT(compilationStats);
+        UNIT_ASSERT(compilationStats->IsFromCache());
+        UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetDurationUs(), 40);
+        UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetDuration().MicroSeconds(), 40);
+        UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetCpuTimeUs(), 50);
+        UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetCpuTime().MicroSeconds(), 50);
+
+        const auto phases = stats.GetQueryPhases();
+        UNIT_ASSERT_VALUES_EQUAL(phases.size(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(phases[0].GetDurationUs(), 60);
+        UNIT_ASSERT_VALUES_EQUAL(phases[0].GetDuration().MicroSeconds(), 60);
+        UNIT_ASSERT_VALUES_EQUAL(phases[0].GetCpuTimeUs(), 70);
+        UNIT_ASSERT_VALUES_EQUAL(phases[0].GetCpuTime().MicroSeconds(), 70);
+        UNIT_ASSERT_VALUES_EQUAL(phases[0].GetAffectedShards(), 80);
+        UNIT_ASSERT(phases[0].IsLiteralPhase());
+
+        const auto& tables = phases[0].GetTableAccess();
+        UNIT_ASSERT_VALUES_EQUAL(tables.size(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetName(), "table");
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetPartitionsCount(), 90);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetReads().GetRows(), 100);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetReads().GetBytes(), 110);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetUpdates().GetRows(), 120);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetUpdates().GetBytes(), 130);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetDeletes().GetRows(), 140);
+        UNIT_ASSERT_VALUES_EQUAL(tables[0].GetDeletes().GetBytes(), 150);
+    }
+
+    Y_UNIT_TEST(KeepsMissingOptionalFieldsEmpty) {
+        const TExecStats stats(Ydb::TableStats::QueryStats{});
+        UNIT_ASSERT(!stats.GetCompilation());
+        UNIT_ASSERT(!stats.GetPlan());
+        UNIT_ASSERT(!stats.GetAst());
+        UNIT_ASSERT(!stats.GetMeta());
+        UNIT_ASSERT(stats.GetQueryPhases().empty());
+    }
+}
diff --git a/ydb/public/sdk/cpp/tests/unit/client/query/ya.make b/ydb/public/sdk/cpp/tests/unit/client/query/ya.make
index 49ea63f30d0..c54a27a4f9c 100644
--- a/ydb/public/sdk/cpp/tests/unit/client/query/ya.make
+++ b/ydb/public/sdk/cpp/tests/unit/client/query/ya.make
@@ -10,14 +10,16 @@ ENDIF()
 FORK_SUBTESTS()
 
 SRCS(
-    client_session_ut.cpp
+   client_session_ut.cpp
+   query_stats_ut.cpp
 )
 
 PEERDIR(
     ydb/public/api/protos
     ydb/public/sdk/cpp/src/library/operation_id
     ydb/public/sdk/cpp/src/client/impl/session
-    ydb/public/sdk/cpp/src/client/query/impl
+    ydb/public/sdk/cpp/src/client/query/impl 
+    ydb/public/sdk/cpp/src/client/query
 )
 
 END()