refactored query stats: implemetation instead interface

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added into query stats the range iterators for iterate over query phases and table access  
+
 ## v3.96.2
 * Fixed broken metric `ydb_go_sdk_ydb_database_sql_conns`
 

internal/query/execute_query.go

@@ -24,7 +24,7 @@ import (
 type executeSettings interface {
 	ExecMode() options.ExecMode
 	StatsMode() options.StatsMode
-	StatsCallback() func(stats stats.QueryStats)
+	StatsCallback() func(stats *stats.QueryStats)
 	TxControl() *query.TransactionControl
 	Syntax() options.Syntax
 	Params() params.Parameters

internal/query/options/execute.go

@@ -32,7 +32,7 @@ type (
 		execMode               ExecMode
 		statsMode              StatsMode
 		resourcePool           string
-		statsCallback          func(queryStats stats.QueryStats)
+		statsCallback          func(queryStats *stats.QueryStats)
 		callOptions            []grpc.CallOption
 		txControl              *tx.Control
 		retryOptions           []retry.Option
@@ -59,7 +59,7 @@ type (
 	syntaxOption    = Syntax
 	statsModeOption struct {
 		mode     StatsMode
-		callback func(stats.QueryStats)
+		callback func(*stats.QueryStats)
 	}
 	execModeOption         = ExecMode
 	responsePartLimitBytes int64
@@ -73,7 +73,7 @@ func (s *executeSettings) RetryOpts() []retry.Option {
 	return s.retryOptions
 }
 
-func (s *executeSettings) StatsCallback() func(stats.QueryStats) {
+func (s *executeSettings) StatsCallback() func(*stats.QueryStats) {
 	return s.statsCallback
 }
 
@@ -225,7 +225,7 @@ func (opt statsModeOption) applyExecuteOption(s *executeSettings) {
 	s.statsCallback = opt.callback
 }
 
-func WithStatsMode(mode StatsMode, callback func(stats.QueryStats)) statsModeOption {
+func WithStatsMode(mode StatsMode, callback func(*stats.QueryStats)) statsModeOption {
 	return statsModeOption{
 		mode:     mode,
 		callback: callback,

internal/query/options/execute_script.go

@@ -34,7 +34,7 @@
 			Query  string
 		}
 		Mode           ExecMode
-		Stats          stats.QueryStats
+		Stats          *stats.QueryStats
 		ResultSetsMeta []struct {
 			Columns []struct {
 				Name string
@@ -78,7 +78,7 @@
 			Query:  pb.GetScriptContent().GetText(),
 		},
 		Mode:  ExecMode(pb.GetExecMode()),
 		Stats: stats.FromQueryStats(pb.GetExecStats()),
 		ResultSetsMeta: func() (
 			resultSetsMeta []struct {
 				Columns []struct {

internal/query/result.go

@@ -38,7 +38,7 @@ type (
 		resultSetIndex int64
 		closed         chan struct{}
 		trace          *trace.Query
-		statsCallback  func(queryStats stats.QueryStats)
+		statsCallback  func(queryStats *stats.QueryStats)
 		onClose        []func()
 		onNextPartErr  []func(err error)
 		onTxMeta       []func(txMeta *Ydb_Query.TransactionMeta)
@@ -93,7 +93,7 @@ func withTrace(t *trace.Query) resultOption {
 	}
 }
 
-func withStatsCallback(callback func(queryStats stats.QueryStats)) resultOption {
+func withStatsCallback(callback func(queryStats *stats.QueryStats)) resultOption {
 	return func(s *streamResult) {
 		s.statsCallback = callback
 	}

internal/query/result_test.go

@@ -1926,8 +1926,8 @@ func TestResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -2287,8 +2287,8 @@ func TestResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -2649,8 +2649,8 @@ func TestResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -2986,8 +2986,8 @@ func TestResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -3358,8 +3358,8 @@ func TestMaterializedResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -3719,8 +3719,8 @@ func TestMaterializedResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -4081,8 +4081,8 @@ func TestMaterializedResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)
@@ -4418,8 +4418,8 @@ func TestMaterializedResultStats(t *testing.T) {
 				},
 			}, nil)
 			stream.EXPECT().Recv().Return(nil, io.EOF)
-			var s stats.QueryStats
-			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats stats.QueryStats) {
+			var s *stats.QueryStats
+			result, err := newResult(ctx, stream, withStatsCallback(func(queryStats *stats.QueryStats) {
 				s = queryStats
 			}))
 			require.NoError(t, err)

internal/stats/query.go

@@ -1,6 +1,7 @@
 package stats
 
 import (
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats"
@@ -19,12 +20,17 @@ type (
 		// NextPhase returns next execution phase within query.
 		// If ok flag is false, then there are no more phases and p is invalid.
 		NextPhase() (p QueryPhase, ok bool)
+
+		// QueryPhases is a range iterator over query phases.
+		QueryPhases() xiter.Seq[QueryPhase]
 	}
 	// QueryPhase holds query execution phase statistics.
 	QueryPhase interface {
 		// NextTableAccess returns next accessed table within query execution phase.
 		// If ok flag is false, then there are no more accessed tables and t is invalid.
 		NextTableAccess() (t *TableAccess, ok bool)
+		// TableAccess is a range iterator over query execution phase's accessed tables.
+		TableAccess() xiter.Seq[*TableAccess]
 		Duration() time.Duration
 		CPUTime() time.Duration
 		AffectedShards() uint64
@@ -86,57 +92,70 @@ func fromOperationStats(pb *Ydb_TableStats.OperationStats) OperationStats {
 	}
 }
 
-func (s *queryStats) ProcessCPUTime() time.Duration {
-	return fromUs(s.pb.GetProcessCpuTimeUs())
+func (stats *queryStats) ProcessCPUTime() time.Duration {
+	return fromUs(stats.pb.GetProcessCpuTimeUs())
 }
 
-func (s *queryStats) Compilation() (c *CompilationStats) {
-	return fromCompilationStats(s.pb.GetCompilation())
+func (stats *queryStats) Compilation() (c *CompilationStats) {
+	return fromCompilationStats(stats.pb.GetCompilation())
 }
 
-func (s *queryStats) QueryPlan() string {
-	return s.pb.GetQueryPlan()
+func (stats *queryStats) QueryPlan() string {
+	return stats.pb.GetQueryPlan()
 }
 
-func (s *queryStats) QueryAST() string {
-	return s.pb.GetQueryAst()
+func (stats *queryStats) QueryAST() string {
+	return stats.pb.GetQueryAst()
 }
 
-func (s *queryStats) TotalCPUTime() time.Duration {
-	return fromUs(s.pb.GetTotalCpuTimeUs())
+func (stats *queryStats) TotalCPUTime() time.Duration {
+	return fromUs(stats.pb.GetTotalCpuTimeUs())
 }
 
-func (s *queryStats) TotalDuration() time.Duration {
-	return fromUs(s.pb.GetTotalDurationUs())
+func (stats *queryStats) TotalDuration() time.Duration {
+	return fromUs(stats.pb.GetTotalDurationUs())
 }
 
 // NextPhase returns next execution phase within query.
 // If ok flag is false, then there are no more phases and p is invalid.
-func (s *queryStats) NextPhase() (p QueryPhase, ok bool) {
-	if s.pos >= len(s.pb.GetQueryPhases()) {
+func (stats *queryStats) NextPhase() (p QueryPhase, ok bool) {
+	if stats.pos >= len(stats.pb.GetQueryPhases()) {
 		return
 	}
-	pb := s.pb.GetQueryPhases()[s.pos]
+	pb := stats.pb.GetQueryPhases()[stats.pos]
 	if pb == nil {
 		return
 	}
-	s.pos++
+	stats.pos++
 
 	return &queryPhase{
 		pb: pb,
 	}, true
 }
 
+func (stats *queryStats) QueryPhases() xiter.Seq[QueryPhase] {
+	return func(yield func(p QueryPhase) bool) {
+		for _, pb := range stats.pb.GetQueryPhases() {
+			cont := yield(&queryPhase{
+				pb: pb,
+			})
+			if !cont {
+				return
+			}
+		}
+	}
+}
+
 // NextTableAccess returns next accessed table within query execution phase.
 //
 // 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(),
@@ -147,20 +166,37 @@ func (queryPhase *queryPhase) NextTableAccess() (t *TableAccess, ok bool) {
 	}, true
 }
 
-func (queryPhase *queryPhase) Duration() time.Duration {
-	return fromUs(queryPhase.pb.GetDurationUs())
+func (phase *queryPhase) TableAccess() 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 {