Support lightweight query statistics modes

README.md

@@ -38,7 +38,8 @@ Flags:
   -i, --instance=STRING            ID of the instance; required for a database
                                    ID ($CLOUDSDK_SPANNER_INSTANCE).
       --database-role=STRING       Database role to assume for all operations.
-      --query-mode="NORMAL"        Query mode.
+      --query-mode="NORMAL"        Query mode: NORMAL, PLAN, PROFILE,
+                                   WITH_PLAN_AND_STATS, or WITH_STATS.
       --priority="unspecified"     Priority for the execute SQL request.
       --format="json"              Output format.
       --redact-rows                Redact result rows from output
@@ -117,6 +118,12 @@ $ execspansql ${DATABASE_ID} --project=${SPANNER_PROJECT} --instance=${SPANNER_I
     --database-role=report_reader --sql='SELECT * FROM Singers'
 ```
 
+### 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<STRING>`; the other modes execute the query and require parameter values.
+
 ### Parameter support
 
 Many Cloud Spanner clients don't support parameter.
@@ -350,4 +357,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.

flags_validation_test.go

@@ -306,6 +306,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"},
@@ -318,6 +328,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"},

main.go

@@ -54,7 +54,7 @@ type opts struct {
 	Project              string        `name:"project" short:"p" env:"CLOUDSDK_CORE_PROJECT" help:"ID of the project; required for a database ID."`
 	Instance             string        `name:"instance" short:"i" env:"CLOUDSDK_SPANNER_INSTANCE" help:"ID of the instance; required for a database ID."`
 	DatabaseRole         string        `name:"database-role" help:"Database role to assume for all operations."`
-	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."`
 	Priority             string        `name:"priority" enum:"high,low,medium,unspecified" default:"unspecified" help:"Priority for the execute SQL request."`
 	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"`
@@ -259,10 +259,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)
 		}
 	}

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

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...)
+}