Added into query stats the range iterators for iterate over query pha…

…ses and table access
ydb-platform · Jan 16, 2025 · b8f2537 · b8f2537
1 parent 7347733
commit b8f2537
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 14 deletions.
diff --git a/internal/stats/query.go b/internal/stats/query.go
@@ -106,7 +106,7 @@ func (s *QueryStats) NextPhase() (p QueryPhase, ok bool) {
 	}, true
 }
 
-func (s *QueryStats) QueryPhases() xiter.Seq[QueryPhase] {
+func (s *QueryStats) RangeQueryPhases() xiter.Seq[QueryPhase] {
 	return func(yield func(p QueryPhase) bool) {
 		for _, pb := range s.pb.GetQueryPhases() {
 			cont := yield(QueryPhase{
@@ -123,12 +123,12 @@ func (s *QueryStats) QueryPhases() xiter.Seq[QueryPhase] {
 //
 // If ok flag is false, then there are no more accessed tables and t is
 // invalid.
-func (QueryPhase *QueryPhase) NextTableAccess() (t *TableAccess, ok bool) {
-	if QueryPhase.pos >= len(QueryPhase.pb.GetTableAccess()) {
+func (phase *QueryPhase) NextTableAccess() (t *TableAccess, ok bool) {
+	if phase.pos >= len(phase.pb.GetTableAccess()) {
 		return
 	}
-	pb := QueryPhase.pb.GetTableAccess()[QueryPhase.pos]
-	QueryPhase.pos++
+	pb := phase.pb.GetTableAccess()[phase.pos]
+	phase.pos++
 
 	return &TableAccess{
 		Name:            pb.GetName(),
@@ -139,20 +139,37 @@ func (QueryPhase *QueryPhase) NextTableAccess() (t *TableAccess, ok bool) {
 	}, true
 }
 
-func (QueryPhase *QueryPhase) Duration() time.Duration {
-	return fromUs(QueryPhase.pb.GetDurationUs())
+func (phase *QueryPhase) RangeTableAccess() xiter.Seq[*TableAccess] {
+	return func(yield func(access *TableAccess) bool) {
+		for _, pb := range phase.pb.GetTableAccess() {
+			cont := yield(&TableAccess{
+				Name:            pb.GetName(),
+				Reads:           fromOperationStats(pb.GetReads()),
+				Updates:         fromOperationStats(pb.GetUpdates()),
+				Deletes:         fromOperationStats(pb.GetDeletes()),
+				PartitionsCount: pb.GetPartitionsCount(),
+			})
+			if !cont {
+				return
+			}
+		}
+	}
+}
+
+func (phase *QueryPhase) Duration() time.Duration {
+	return fromUs(phase.pb.GetDurationUs())
 }
 
-func (QueryPhase *QueryPhase) CPUTime() time.Duration {
-	return fromUs(QueryPhase.pb.GetCpuTimeUs())
+func (phase *QueryPhase) CPUTime() time.Duration {
+	return fromUs(phase.pb.GetCpuTimeUs())
 }
 
-func (QueryPhase *QueryPhase) AffectedShards() uint64 {
-	return QueryPhase.pb.GetAffectedShards()
+func (phase *QueryPhase) AffectedShards() uint64 {
+	return phase.pb.GetAffectedShards()
 }
 
-func (QueryPhase *QueryPhase) IsLiteralPhase() bool {
-	return QueryPhase.pb.GetLiteralPhase()
+func (phase *QueryPhase) IsLiteralPhase() bool {
+	return phase.pb.GetLiteralPhase()
 }
 
 func FromQueryStats(pb *Ydb_TableStats.QueryStats) *QueryStats {

diff --git a/internal/stats/query_go1.23_test.go b/internal/stats/query_go1.23_test.go
@@ -16,12 +16,45 @@ func TestIterateOverQueryPhases(t *testing.T) {
 			QueryPhases: []*Ydb_TableStats.QueryPhaseStats{
 				{
 					DurationUs: 1,
+					TableAccess: []*Ydb_TableStats.TableAccessStats{
+						{
+							Name: "a",
+						},
+						{
+							Name: "b",
+						},
+						{
+							Name: "c",
+						},
+					},
 				},
 				{
 					DurationUs: 2,
+					TableAccess: []*Ydb_TableStats.TableAccessStats{
+						{
+							Name: "d",
+						},
+						{
+							Name: "e",
+						},
+						{
+							Name: "f",
+						},
+					},
 				},
 				{
 					DurationUs: 3,
+					TableAccess: []*Ydb_TableStats.TableAccessStats{
+						{
+							Name: "g",
+						},
+						{
+							Name: "h",
+						},
+						{
+							Name: "i",
+						},
+					},
 				},
 			},
 		},
@@ -30,23 +63,37 @@ func TestIterateOverQueryPhases(t *testing.T) {
 		for i := range make([]struct{}, 3) {
 			t.Run(fmt.Sprintf("Pass#%d", i), func(t *testing.T) {
 				durations := make([]uint64, 0, 3)
-				for phase := range s.QueryPhases() {
+				tables := make([]string, 0, 9)
+				for phase := range s.RangeQueryPhases() {
 					durations = append(durations, phase.pb.GetDurationUs())
+					for access := range phase.RangeTableAccess() {
+						tables = append(tables, access.Name)
+					}
 				}
 				require.Equal(t, []uint64{1, 2, 3}, durations)
+				require.Equal(t, []string{"a", "b", "c", "d", "e", "f", "g", "h", "i"}, tables)
 			})
 		}
 	})
 	t.Run("MutableIteration", func(t *testing.T) {
 		durations := make([]uint64, 0, 3)
+		tables := make([]string, 0, 9)
 		for {
 			phase, ok := s.NextPhase()
 			if !ok {
 				break
 			}
 			durations = append(durations, phase.pb.GetDurationUs())
+			for {
+				access, ok := phase.NextTableAccess()
+				if !ok {
+					break
+				}
+				tables = append(tables, access.Name)
+			}
 		}
 		require.Equal(t, []uint64{1, 2, 3}, durations)
+		require.Equal(t, []string{"a", "b", "c", "d", "e", "f", "g", "h", "i"}, tables)
 		require.Equal(t, 3, s.pos)
 
 		_, ok := s.NextPhase()