From 7e7142bda89aad1fb5072515529efb35f6a55eff Mon Sep 17 00:00:00 2001 From: apstndb <803393+apstndb@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:31:17 +0900 Subject: [PATCH] Support lightweight query statistics modes --- README.md | 11 ++- flags_validation_test.go | 22 +++++ main.go | 9 +- pdml_query_mode_test.go | 2 +- query_stats_modes_test.go | 168 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 205 insertions(+), 7 deletions(-) create mode 100644 query_stats_modes_test.go diff --git a/README.md b/README.md index 9d52320..9dcc702 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,8 @@ Flags: -p, --project=STRING ID of the project ($CLOUDSDK_CORE_PROJECT). -i, --instance=STRING ID of the instance ($CLOUDSDK_SPANNER_INSTANCE). - --query-mode="NORMAL" Query mode. + --query-mode="NORMAL" Query mode: NORMAL, PLAN, PROFILE, + WITH_PLAN_AND_STATS, or WITH_STATS. --format="json" Output format. --redact-rows Redact result rows from output -c, --compact-output Compact JSON output (--compact-output of jq) @@ -95,6 +96,12 @@ $ docker run --rm -t -v "${HOME}/.config/gcloud/application_default_credentials. There are examples omitting some required options. +### Query modes + +`--query-mode` matches gcloud's query-stat modes. `NORMAL` returns result rows only; `PLAN` returns the plan without rows or execution statistics; `PROFILE` returns rows, a plan, overall statistics, and operator-level statistics. `WITH_PLAN_AND_STATS` returns rows, a plan, and overall statistics without operator-level statistics. `WITH_STATS` returns rows and overall statistics without a plan or operator-level statistics. + +Only `PLAN` accepts bare parameter type expressions such as `ARRAY`; the other modes execute the query and require parameter values. + ### Parameter support Many Cloud Spanner clients don't support parameter. @@ -324,4 +331,4 @@ exit status 1 * `--format=experimental_csv` does not run the jq pipeline; `--filter`, `--filter-file`, `--raw-output`, `--compact-output`, and `--jq-input-mode=lazy` are rejected. * `--raw-output` and `--compact-output` are supported only when `--format=json`. -* `--query-mode=PLAN` and `--query-mode=PROFILE` cannot be combined with `--enable-partitioned-dml`. The Partitioned DML client path ignores query mode and would execute writes. +* Non-`NORMAL` query modes (`PLAN`, `PROFILE`, `WITH_PLAN_AND_STATS`, and `WITH_STATS`) cannot be combined with `--enable-partitioned-dml`. The Partitioned DML client path ignores query mode and would execute writes. diff --git a/flags_validation_test.go b/flags_validation_test.go index 641718d..8873d72 100644 --- a/flags_validation_test.go +++ b/flags_validation_test.go @@ -298,6 +298,16 @@ func TestValidateExecutionOptions(t *testing.T) { o: opts{QueryMode: "PROFILE"}, mode: readWrite{}, }, + { + name: "with_plan_and_stats_allows_read_write_dml", + o: opts{QueryMode: "WITH_PLAN_AND_STATS"}, + mode: readWrite{}, + }, + { + name: "with_stats_allows_read_write_dml", + o: opts{QueryMode: "WITH_STATS"}, + mode: readWrite{}, + }, { name: "plan_rejects_partitioned_dml", o: opts{EnablePartitionedDML: true, QueryMode: "PLAN"}, @@ -310,6 +320,18 @@ func TestValidateExecutionOptions(t *testing.T) { mode: partitionedDML{}, err: "--query-mode=PROFILE cannot be combined with --enable-partitioned-dml", }, + { + name: "with_plan_and_stats_rejects_partitioned_dml", + o: opts{EnablePartitionedDML: true, QueryMode: "WITH_PLAN_AND_STATS"}, + mode: partitionedDML{}, + err: "--query-mode=WITH_PLAN_AND_STATS cannot be combined with --enable-partitioned-dml", + }, + { + name: "with_stats_rejects_partitioned_dml", + o: opts{EnablePartitionedDML: true, QueryMode: "WITH_STATS"}, + mode: partitionedDML{}, + err: "--query-mode=WITH_STATS cannot be combined with --enable-partitioned-dml", + }, { name: "plan_rejects_partitioned_dml_json", o: opts{EnablePartitionedDML: true, QueryMode: "PLAN", Format: "json"}, diff --git a/main.go b/main.go index ccdeedc..e9bacc9 100644 --- a/main.go +++ b/main.go @@ -53,7 +53,7 @@ type opts struct { SqlFile string `name:"sql-file" xor:"sql" required:"" help:"File name contains SQL query; exclusive with --sql"` Project string `name:"project" short:"p" env:"CLOUDSDK_CORE_PROJECT" required:"" help:"ID of the project."` Instance string `name:"instance" short:"i" env:"CLOUDSDK_SPANNER_INSTANCE" required:"" help:"ID of the instance."` - QueryMode string `name:"query-mode" enum:"NORMAL,PLAN,PROFILE" default:"NORMAL" help:"Query mode."` + QueryMode string `name:"query-mode" enum:"NORMAL,PLAN,PROFILE,WITH_PLAN_AND_STATS,WITH_STATS" default:"NORMAL" help:"Query mode: NORMAL, PLAN, PROFILE, WITH_PLAN_AND_STATS, or WITH_STATS."` Format string `name:"format" enum:"json,yaml,experimental_csv" default:"json" help:"Output format."` RedactRows bool `name:"redact-rows" help:"Redact result rows from output"` CompactOutput bool `name:"compact-output" short:"c" help:"Compact JSON output (--compact-output of jq)"` @@ -219,10 +219,11 @@ func validateExecutionOptions(o opts, mode queryMode) error { if _, ok := mode.(partitionedDML); !ok { return fmt.Errorf("--enable-partitioned-dml can only be used with DML statements") } - // PartitionedUpdateWithOptions does not copy QueryOptions.Mode, so PLAN - // and PROFILE would execute writes instead of returning a plan or profile. + // PartitionedUpdateWithOptions does not copy QueryOptions.Mode. Every + // non-NORMAL query mode would therefore execute writes instead of + // returning its requested plan and/or statistics. switch o.QueryMode { - case "PLAN", "PROFILE": + case "PLAN", "PROFILE", "WITH_PLAN_AND_STATS", "WITH_STATS": return fmt.Errorf("--query-mode=%s cannot be combined with --enable-partitioned-dml", o.QueryMode) } } diff --git a/pdml_query_mode_test.go b/pdml_query_mode_test.go index 0da2a05..dc4f1e1 100644 --- a/pdml_query_mode_test.go +++ b/pdml_query_mode_test.go @@ -39,7 +39,7 @@ func TestPartitionedDMLQueryMode(t *testing.T) { "--enable-partitioned-dml", } - for _, queryMode := range []string{"PLAN", "PROFILE"} { + for _, queryMode := range []string{"PLAN", "PROFILE", "WITH_PLAN_AND_STATS", "WITH_STATS"} { for _, format := range []string{"json", "yaml", "experimental_csv"} { queryMode := queryMode format := format diff --git a/query_stats_modes_test.go b/query_stats_modes_test.go new file mode 100644 index 0000000..fb0822e --- /dev/null +++ b/query_stats_modes_test.go @@ -0,0 +1,168 @@ +package main + +import ( + "context" + "net" + "strings" + "sync" + "testing" + + "cloud.google.com/go/spanner" + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/structpb" +) + +func TestAdditionalQueryStatsModesRejectPartitionedDMLBeforeClient(t *testing.T) { + for _, queryMode := range []string{"WITH_PLAN_AND_STATS", "WITH_STATS"} { + queryMode := queryMode + t.Run(queryMode, func(t *testing.T) { + err := runMain(t, []string{ + "database", "--project", "unused-project", "--instance", "unused-instance", + "--sql", "UPDATE T SET V=1", "--enable-partitioned-dml", "--query-mode", queryMode, + }) + want := "--query-mode=" + queryMode + " cannot be combined with --enable-partitioned-dml" + if err == nil || !strings.Contains(err.Error(), want) { + t.Fatalf("_main() error = %v, want %q", err, want) + } + }) + } +} + +func TestQueryStatsModesPreserveDMLCounts(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mode sppb.ExecuteSqlRequest_QueryMode + wantDMLRowCount bool + }{ + {name: "normal", mode: sppb.ExecuteSqlRequest_NORMAL, wantDMLRowCount: true}, + {name: "plan", mode: sppb.ExecuteSqlRequest_PLAN}, + {name: "profile", mode: sppb.ExecuteSqlRequest_PROFILE, wantDMLRowCount: true}, + {name: "with_plan_and_stats", mode: sppb.ExecuteSqlRequest_WITH_PLAN_AND_STATS, wantDMLRowCount: true}, + {name: "with_stats", mode: sppb.ExecuteSqlRequest_WITH_STATS, wantDMLRowCount: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := spanner.QueryOptions{Mode: tt.mode.Enum()} + if got := dmlRowCountForMode(readWrite{}, opts); got != tt.wantDMLRowCount { + t.Errorf("dmlRowCountForMode() = %v, want %v", got, tt.wantDMLRowCount) + } + }) + } +} + +func TestAdditionalQueryStatsModesReachSpannerAndProduceStats(t *testing.T) { + server := &queryStatsModeServer{} + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + grpcServer := grpc.NewServer() + sppb.RegisterSpannerServer(grpcServer, server) + go func() { _ = grpcServer.Serve(listener) }() + t.Cleanup(func() { + grpcServer.Stop() + _ = listener.Close() + }) + t.Setenv("SPANNER_EMULATOR_HOST", listener.Addr().String()) + + for _, queryMode := range []string{"WITH_PLAN_AND_STATS", "WITH_STATS"} { + for _, tc := range []struct { + name string + format string + lazy bool + }{ + {name: "json_eager", format: "json"}, + {name: "yaml_eager", format: "yaml"}, + {name: "json_lazy", format: "json", lazy: true}, + {name: "yaml_lazy", format: "yaml", lazy: true}, + } { + queryMode := queryMode + tc := tc + t.Run(queryMode+"_"+tc.name, func(t *testing.T) { + server.resetModes() + args := []string{ + "database", "--project", "project", "--instance", "instance", "--sql", "SELECT 'value'", + "--query-mode", queryMode, "--format", tc.format, "--timeout", "5s", + } + if tc.lazy { + args = append(args, "--jq-input-mode", "lazy", "--filter", ".stats.queryStats.mode") + } + out, err := captureStdout(t, func() error { return runMain(t, args) }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(out, queryMode) { + t.Fatalf("output = %q, want query stats containing %q", out, queryMode) + } + if !tc.lazy { + if !strings.Contains(out, "value") { + t.Fatalf("eager output = %q, want row value", out) + } + if queryMode == "WITH_PLAN_AND_STATS" && !strings.Contains(out, "Fake Scan") { + t.Fatalf("WITH_PLAN_AND_STATS output = %q, want fake query plan", out) + } + if queryMode == "WITH_STATS" && strings.Contains(out, "Fake Scan") { + t.Fatalf("WITH_STATS output = %q, got unexpected query plan", out) + } + } + modes := server.modes() + if len(modes) != 1 || modes[0].String() != queryMode { + t.Fatalf("received query modes = %v, want [%s]", modes, queryMode) + } + }) + } + } +} + +type queryStatsModeServer struct { + sppb.UnimplementedSpannerServer + + mu sync.Mutex + receivedModes []sppb.ExecuteSqlRequest_QueryMode +} + +func (s *queryStatsModeServer) CreateSession(_ context.Context, req *sppb.CreateSessionRequest) (*sppb.Session, error) { + return &sppb.Session{Name: req.GetDatabase() + "/sessions/test"}, nil +} + +func (s *queryStatsModeServer) ExecuteStreamingSql(req *sppb.ExecuteSqlRequest, stream sppb.Spanner_ExecuteStreamingSqlServer) error { + s.mu.Lock() + s.receivedModes = append(s.receivedModes, req.GetQueryMode()) + s.mu.Unlock() + + if err := stream.Send(&sppb.PartialResultSet{ + Metadata: &sppb.ResultSetMetadata{RowType: &sppb.StructType{Fields: []*sppb.StructType_Field{{ + Name: "value", + Type: &sppb.Type{Code: sppb.TypeCode_STRING}, + }}}}, + Values: []*structpb.Value{structpb.NewStringValue("value")}, + }); err != nil { + return err + } + + stats := &sppb.ResultSetStats{ + QueryStats: &structpb.Struct{Fields: map[string]*structpb.Value{ + "mode": structpb.NewStringValue(req.GetQueryMode().String()), + }}, + } + if req.GetQueryMode() == sppb.ExecuteSqlRequest_WITH_PLAN_AND_STATS { + stats.QueryPlan = &sppb.QueryPlan{PlanNodes: []*sppb.PlanNode{{DisplayName: "Fake Scan"}}} + } + return stream.Send(&sppb.PartialResultSet{Stats: stats}) +} + +func (s *queryStatsModeServer) resetModes() { + s.mu.Lock() + defer s.mu.Unlock() + s.receivedModes = nil +} + +func (s *queryStatsModeServer) modes() []sppb.ExecuteSqlRequest_QueryMode { + s.mu.Lock() + defer s.mu.Unlock() + return append([]sppb.ExecuteSqlRequest_QueryMode(nil), s.receivedModes...) +}