* Added options.WithExecuteScanQueryStats option

* Added to query stats plan and AST

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,6 @@
+* Added `options.WithExecuteScanQueryStats` option
+* Added to query stats plan and AST
+* Changed behaviour of `result.Stats()` (if query result have no stats - returns `nil`) 
 * Added context cancel with specific error
 
 ## v3.26.10

internal/table/scanner/result.go

@@ -149,13 +149,18 @@ func (r *baseResult) CurrentResultSet() result.Set {
 
 // Stats returns query execution queryStats.
 func (r *baseResult) Stats() stats.QueryStats {
-	var s queryStats
 	r.statsMtx.RLock()
-	s.stats = r.stats
+	stats := r.stats
 	r.statsMtx.RUnlock()
-	s.processCPUTime = time.Microsecond * time.Duration(s.stats.GetProcessCpuTimeUs())
-	s.pos = 0
-	return &s
+
+	if stats == nil {
+		return nil
+	}
+
+	return &queryStats{
+		stats:          stats,
+		processCPUTime: time.Microsecond * time.Duration(stats.GetProcessCpuTimeUs()),
+	}
 }
 
 // Close closes the result, preventing further iteration.

internal/table/scanner/stats.go

@@ -16,31 +16,32 @@ type queryStats struct {
 }
 
 func (s *queryStats) ProcessCPUTime() time.Duration {
-	if s == nil {
-		return 0
-	}
 	return s.processCPUTime
 }
 
 func (s *queryStats) Compilation() (c *stats.CompilationStats) {
-	if s == nil || s.stats == nil || s.stats.Compilation == nil {
+	if s.stats == nil || s.stats.Compilation == nil {
 		return nil
 	}
-	x := s.stats.Compilation
-	if x == nil {
-		return
-	}
 	return &stats.CompilationStats{
-		FromCache: x.FromCache,
-		Duration:  time.Microsecond * time.Duration(x.DurationUs),
-		CPUTime:   time.Microsecond * time.Duration(x.CpuTimeUs),
+		FromCache: s.stats.Compilation.FromCache,
+		Duration:  time.Microsecond * time.Duration(s.stats.Compilation.DurationUs),
+		CPUTime:   time.Microsecond * time.Duration(s.stats.Compilation.CpuTimeUs),
 	}
 }
 
+func (s *queryStats) QueryPlan() string {
+	return s.stats.GetQueryPlan()
+}
+
+func (s *queryStats) QueryAST() string {
+	return s.stats.GetQueryAst()
+}
+
 // 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 stats.QueryPhase, ok bool) {
-	if s.stats == nil || s.pos >= len(s.stats.QueryPhases) {
+	if s.pos >= len(s.stats.QueryPhases) {
 		return
 	}
 	x := s.stats.QueryPhases[s.pos]
@@ -54,6 +55,7 @@ func (s *queryStats) NextPhase() (p stats.QueryPhase, ok bool) {
 		duration:       time.Microsecond * time.Duration(x.DurationUs),
 		cpuTime:        time.Microsecond * time.Duration(x.CpuTimeUs),
 		affectedShards: x.AffectedShards,
+		literalPhase:   x.LiteralPhase,
 	}, true
 }
 
@@ -64,6 +66,7 @@ type queryPhase struct {
 	affectedShards uint64
 	tables         []*Ydb_TableStats.TableAccessStats
 	pos            int
+	literalPhase   bool
 }
 
 // NextTableAccess returns next accessed table within query execution phase.
@@ -96,6 +99,10 @@ func (q *queryPhase) AffectedShards() uint64 {
 	return q.affectedShards
 }
 
+func (q *queryPhase) IsLiteralPhase() bool {
+	return q.literalPhase
+}
+
 func initOperationStats(x *Ydb_TableStats.OperationStats) stats.OperationStats {
 	if x == nil {
 		return stats.OperationStats{}

table/options/options.go

@@ -583,12 +583,42 @@ type (
 	ExecuteScanQueryOption func(*ExecuteScanQueryDesc)
 )
 
+// WithExecuteScanQueryMode defines scan query mode: execute or explain
 func WithExecuteScanQueryMode(m ExecuteScanQueryRequestMode) ExecuteScanQueryOption {
 	return func(desc *ExecuteScanQueryDesc) {
 		desc.Mode = m.toYDB()
 	}
 }
 
+// ExecuteScanQueryStatsType specified scan query mode
+type ExecuteScanQueryStatsType uint32
+
+const (
+	ExecuteScanQueryStatsTypeNone = iota
+	ExecuteScanQueryStatsTypeBasic
+	ExecuteScanQueryStatsTypeFull
+)
+
+func (stats ExecuteScanQueryStatsType) toYDB() Ydb_Table.QueryStatsCollection_Mode {
+	switch stats {
+	case ExecuteScanQueryStatsTypeNone:
+		return Ydb_Table.QueryStatsCollection_STATS_COLLECTION_NONE
+	case ExecuteScanQueryStatsTypeBasic:
+		return Ydb_Table.QueryStatsCollection_STATS_COLLECTION_BASIC
+	case ExecuteScanQueryStatsTypeFull:
+		return Ydb_Table.QueryStatsCollection_STATS_COLLECTION_FULL
+	default:
+		return Ydb_Table.QueryStatsCollection_STATS_COLLECTION_UNSPECIFIED
+	}
+}
+
+// WithExecuteScanQueryStats defines query statistics mode
+func WithExecuteScanQueryStats(stats ExecuteScanQueryStatsType) ExecuteScanQueryOption {
+	return func(desc *ExecuteScanQueryDesc) {
+		desc.CollectStats = stats.toYDB()
+	}
+}
+
 // Read table options
 type (
 	ReadTableDesc   Ydb_Table.ReadTableRequest

table/result/result.go

@@ -106,6 +106,8 @@ type result interface {
 	ScanNamed(namedValues ...named.Value) error
 
 	// Stats returns query execution QueryStats.
+	//
+	// If query result have no stats - returns nil
 	Stats() (s stats.QueryStats)
 
 	// Err return scanner error

table/stats/stats.go

@@ -10,12 +10,15 @@ type QueryPhase interface {
 	Duration() time.Duration
 	CPUTime() time.Duration
 	AffectedShards() uint64
+	IsLiteralPhase() bool
 }
 
 // QueryStats holds query execution statistics.
 type QueryStats interface {
 	ProcessCPUTime() time.Duration
 	Compilation() (c *CompilationStats)
+	QueryPlan() string
+	QueryAST() string
 	// 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)

table/table_e2e_test.go

@@ -559,6 +559,7 @@ func TestTable(t *testing.T) {
 		func(ctx context.Context, s table.Session) (err error) {
 			res, err := s.StreamExecuteScanQuery(
 				ctx, `SELECT val FROM stream_query;`, table.NewQueryParameters(),
+				options.WithExecuteScanQueryStats(options.ExecuteScanQueryStatsTypeFull),
 			)
 			if err != nil {
 				return err
@@ -579,6 +580,21 @@ func TestTable(t *testing.T) {
 					}
 					checkSum += uint64(*val)
 				}
+				if stats := res.Stats(); stats != nil {
+					t.Logf(" --- query stats: compilation: %v, process CPU time: %v, affected shards: %v\n",
+						stats.Compilation(),
+						stats.ProcessCPUTime(),
+						func() (count uint64) {
+							for {
+								phase, ok := stats.NextPhase()
+								if !ok {
+									return
+								}
+								count += phase.AffectedShards()
+							}
+						}(),
+					)
+				}
 			}
 			if rowsCount != upsertRowsCount {
 				return fmt.Errorf("wrong rows count: %v, exp: %v", rowsCount, upsertRowsCount)
@@ -695,27 +711,30 @@ func streamReadTable(ctx context.Context, t *testing.T, c table.Client, tableAbs
 			if err = res.Err(); err != nil {
 				return err
 			}
-			stats := res.Stats()
-			for i := 0; ; i++ {
-				phase, ok := stats.NextPhase()
-				if !ok {
-					break
-				}
-				t.Logf(
-					"# phase #%d: took %s\n",
-					i, phase.Duration(),
-				)
-				for {
-					tbl, ok := phase.NextTableAccess()
+
+			if stats := res.Stats(); stats != nil {
+				for i := 0; ; i++ {
+					phase, ok := stats.NextPhase()
 					if !ok {
 						break
 					}
 					t.Logf(
-						"#  accessed %s: read=(%drows, %dbytes)\n",
-						tbl.Name, tbl.Reads.Rows, tbl.Reads.Bytes,
+						"# phase #%d: took %s\n",
+						i, phase.Duration(),
 					)
+					for {
+						tbl, ok := phase.NextTableAccess()
+						if !ok {
+							break
+						}
+						t.Logf(
+							"#  accessed %s: read=(%drows, %dbytes)\n",
+							tbl.Name, tbl.Reads.Rows, tbl.Reads.Bytes,
+						)
+					}
 				}
 			}
+
 			return nil
 		},
 	)