diff --git a/internal/stats/query.go b/internal/stats/query.go index 3b35beefe..90070a40a 100644 --- a/internal/stats/query.go +++ b/internal/stats/query.go @@ -4,6 +4,8 @@ import ( "time" "github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats" + + "github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter" ) type ( @@ -104,6 +106,19 @@ func (s *QueryStats) NextPhase() (p QueryPhase, ok bool) { }, true } +func (s *QueryStats) QueryPhases() xiter.Seq[QueryPhase] { + return func(yield func(p QueryPhase) bool) { + for _, pb := range s.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 diff --git a/internal/stats/query_go1.23_test.go b/internal/stats/query_go1.23_test.go new file mode 100644 index 000000000..bc00a8011 --- /dev/null +++ b/internal/stats/query_go1.23_test.go @@ -0,0 +1,55 @@ +//go:build go1.23 + +package stats + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/require" + "github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats" +) + +func TestIterateOverQueryPhases(t *testing.T) { + s := &QueryStats{ + pb: &Ydb_TableStats.QueryStats{ + QueryPhases: []*Ydb_TableStats.QueryPhaseStats{ + { + DurationUs: 1, + }, + { + DurationUs: 2, + }, + { + DurationUs: 3, + }, + }, + }, + } + t.Run("ImmutableIteration", func(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() { + durations = append(durations, phase.pb.GetDurationUs()) + } + require.Equal(t, []uint64{1, 2, 3}, durations) + }) + } + }) + t.Run("MutableIteration", func(t *testing.T) { + durations := make([]uint64, 0, 3) + for { + phase, ok := s.NextPhase() + if !ok { + break + } + durations = append(durations, phase.pb.GetDurationUs()) + } + require.Equal(t, []uint64{1, 2, 3}, durations) + require.Equal(t, 3, s.pos) + + _, ok := s.NextPhase() + require.False(t, ok) + }) +} diff --git a/internal/xiter/xiter.go b/internal/xiter/xiter.go index d2fc0ab5d..1a8d5d521 100644 --- a/internal/xiter/xiter.go +++ b/internal/xiter/xiter.go @@ -2,4 +2,7 @@ package xiter -type Seq2[K, V any] func(yield func(K, V) bool) +type ( + Seq[T any] func(yield func(T) bool) + Seq2[K, V any] func(yield func(K, V) bool) +) diff --git a/internal/xiter/xiter_go1.23.go b/internal/xiter/xiter_go1.23.go index 77a2b2adc..12253f9cf 100644 --- a/internal/xiter/xiter_go1.23.go +++ b/internal/xiter/xiter_go1.23.go @@ -6,4 +6,7 @@ import ( "iter" ) -type Seq2[K, V any] iter.Seq2[K, V] +type ( + Seq[T any] iter.Seq[T] + Seq2[K, V any] iter.Seq2[K, V] +)