Split query plan and row output into separate sinks

README.md

@@ -13,6 +13,7 @@ Yet another `gcloud spanner databases execute-sql` replacement for better compos
 * Embedded jq
 * Configurable gRPC logging (`off`, `metadata`, `payload` with payload caveat)
 * (Experimental) CSV output
+* Split query-plan and row output (`--plan-output`)
 * (Experimental) Check whether the query can be executed as a partition query or not.
 
 This tool is still pre-release quality and none of guarantees.
@@ -41,7 +42,18 @@ Flags:
       --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.
+      --format="json"              Output format of the primary document.
+  -o, --output="-"                 Destination of the primary document. Use -
+                                   for stdout; /dev/stdout and /dev/stderr are
+                                   mapped in-process.
+      --plan-output=STRING         Write the query-plan artifact here and strip
+                                   stats.queryPlan from the primary document.
+                                   Enables split mode.
+      --plan-format=STRING         Format of the plan artifact: json or yaml.
+                                   Defaults to --format when that is json or
+                                   yaml, otherwise json. Requires --plan-output.
+      --discard-results            Do not write the primary document
+                                   (plan-only). Requires --plan-output.
       --redact-rows                Redact result rows from output
   -c, --compact-output             Compact JSON output (--compact-output of jq)
       --filter=STRING              jq filter
@@ -125,6 +137,53 @@ $ execspansql ${DATABASE_ID} --project=${SPANNER_PROJECT} --instance=${SPANNER_I
 
 Only `PLAN` accepts bare parameter type expressions such as `ARRAY<STRING>`; the other modes execute the query and require parameter values.
 
+### Split plan and row output
+
+Setting `--plan-output` switches from the default combined document to split mode. The default (no `--plan-output`) stays a single combined `ResultSet` (or CSV of rows) on stdout, byte-for-byte identical to previous versions.
+
+| Flag | Default | Meaning |
+|------|---------|---------|
+| `--output PATH` (`-o`) | `-` | Destination of the primary document (metadata, rows, `stats` without `queryPlan`). |
+| `--plan-output PATH` | unset | Enables split mode: write the plan artifact here and remove `stats.queryPlan` from the primary document. |
+| `--plan-format json\|yaml` | follows `--format` when that is `json` or `yaml`, otherwise `json` | Format of the plan artifact. |
+| `--discard-results` | off | Do not write the primary document (plan-only). Requires `--plan-output`. |
+
+`--redact-rows` is independent of `--discard-results`: redact still emits metadata and a CSV header; discard writes no primary bytes at all.
+
+Path conventions for both `--output` and `--plan-output`:
+
+- `-` means stdout.
+- `/dev/stdout` and `/dev/stderr` are recognized literally and mapped to stdout/stderr in-process (so they work on Windows and take part in collision checks). `--plan-output=/dev/stderr` is the supported spelling for "plan on the terminal while rows go down the pipe".
+- Any other value is a regular file. Files are written to a sibling temp file (mode `0600`) and renamed into place after the query (and transaction) succeeds, so a failing query leaves an existing target intact. Overwriting an existing target is allowed. Two files plus stdout are not a transaction: if publishing the second file fails, the command reports which outputs completed and exits non-zero. SQL is never replayed because an output failed.
+
+In split mode the two destinations must differ. Both on stdout (any spelling) is rejected. `--plan-output` requires `--query-mode=PLAN`, `PROFILE`, or `WITH_PLAN_AND_STATS` (never upgraded from `NORMAL` or `WITH_STATS`). It is incompatible with `--try-partition-query` and `--enable-partitioned-dml`.
+
+Document contents:
+
+- Primary document: the current `ResultSet` with `stats.queryPlan` removed. `stats.queryStats` and `stats.rowCount*` stay. CSV primary output is unchanged (rows only).
+- Plan artifact: a `ResultSet` envelope without `rows` — `metadata` (for `rowType`) plus the full `stats` (`queryPlan`, `queryStats`, `rowCount*`). jq flags apply only to the primary document; the plan is never filtered.
+
+Split mode disables jq early stop: remaining rows are drained so the final plan/stats can be captured, at the same server cost as reading everything. Rows drained only for the plan are not retained. `--jq-input-mode=lazy` still caches rows that jq actually consumed.
+
+If output or rendering fails after a committed DML statement, the process exits non-zero and says so. That failure is not a rollback and the SQL is not replayed.
+
+```
+$ execspansql ${DATABASE_ID} --query-mode=PROFILE --format=experimental_csv \
+    --output=- --plan-output=/dev/stderr \
+    --sql='SELECT * FROM Singers'
+```
+
+```
+$ execspansql ${DATABASE_ID} --query-mode=PROFILE \
+    --output=rows.yaml --format=yaml --plan-output=plan.json --plan-format=json \
+    --sql='SELECT * FROM Singers'
+```
+
+```
+$ execspansql ${DATABASE_ID} --query-mode=PROFILE --discard-results \
+    --plan-output=plan.json --sql='SELECT * FROM Singers'
+```
+
 ### Parameter support
 
 Many Cloud Spanner clients don't support parameter.
@@ -192,7 +251,7 @@ execspansql can process output using embedded [wader/gojq](https://github.com/wa
 
 In `lazy` mode, `metadata` is populated after the first row is read from Spanner (or after a zero-row result). Prefer `.rows[]` to stream rows. Bare `.rows` is a lazy iterator: reuse it in one object literal (for example `{a: .rows, b: .rows}`) may not duplicate rows because jq can evaluate the subexpression once; use `{a: [.rows[]], b: [.rows[]]}` when you need two row arrays. After `.stats` drains the iterator, captured `.rows` values replay from materialized rows.
 
-`--jq-input-mode=lazy` emits rows incrementally, but rows are cached internally after first materialization and reused, so it is not a strict constant-memory mode for large result sets.
+`--jq-input-mode=lazy` emits rows incrementally, but rows are cached internally after first materialization and reused, so it is not a strict constant-memory mode for large result sets. Split mode (`--plan-output`) disables jq early stop: remaining rows are still drained so the plan artifact can be written.
 
 Output expands top-level `gojq.Iter` to one JSON/YAML document per row (JSONL-style). Nested `Iter` values inside objects are expanded to arrays on encode.
 
@@ -362,3 +421,5 @@ 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`.
 * 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.
+* `--plan-output` requires a plan-producing query mode and cannot be combined with `--try-partition-query` or `--enable-partitioned-dml`.
+* Split mode disables jq early stop so the plan artifact can be captured after the last `PartialResultSet`.

flags_validation_test.go

@@ -373,6 +373,63 @@ func TestValidateExecutionOptions(t *testing.T) {
 			o:    opts{EnablePartitionedDML: true, QueryMode: "NORMAL", Format: "experimental_csv"},
 			mode: partitionedDML{},
 		},
+		{
+			name: "plan_output_allows_profile",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "PROFILE"},
+			mode: single{spanner.StrongRead()},
+		},
+		{
+			name: "plan_output_allows_plan_mode",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "PLAN"},
+			mode: single{spanner.StrongRead()},
+		},
+		{
+			name: "plan_output_allows_with_plan_and_stats",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "WITH_PLAN_AND_STATS"},
+			mode: single{spanner.StrongRead()},
+		},
+		{
+			name: "plan_output_rejects_normal",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "NORMAL"},
+			mode: single{spanner.StrongRead()},
+			err:  "--plan-output requires --query-mode=PLAN, PROFILE, or WITH_PLAN_AND_STATS",
+		},
+		{
+			name: "plan_output_rejects_with_stats",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "WITH_STATS"},
+			mode: single{spanner.StrongRead()},
+			err:  "--plan-output requires --query-mode=PLAN, PROFILE, or WITH_PLAN_AND_STATS",
+		},
+		{
+			name: "plan_format_requires_plan_output",
+			o:    opts{PlanFormat: "json", QueryMode: "PROFILE"},
+			mode: single{spanner.StrongRead()},
+			err:  "--plan-format requires --plan-output",
+		},
+		{
+			name: "discard_results_requires_plan_output",
+			o:    opts{DiscardResults: true, QueryMode: "PROFILE"},
+			mode: single{spanner.StrongRead()},
+			err:  "--discard-results requires --plan-output",
+		},
+		{
+			name: "plan_output_rejects_try_partition_query",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "PROFILE", TryPartitionQuery: true},
+			mode: single{spanner.StrongRead()},
+			err:  "--plan-output cannot be combined with --try-partition-query",
+		},
+		{
+			name: "plan_output_rejects_partitioned_dml",
+			o:    opts{PlanOutput: "plan.json", QueryMode: "PROFILE", EnablePartitionedDML: true},
+			mode: partitionedDML{},
+			err:  "--plan-output cannot be combined with --enable-partitioned-dml",
+		},
+		{
+			name: "plan_format_rejects_unknown",
+			o:    opts{PlanOutput: "plan.json", PlanFormat: "text", QueryMode: "PROFILE"},
+			mode: single{spanner.StrongRead()},
+			err:  "--plan-format must be json or yaml",
+		},
 	}
 
 	for _, tt := range tests {

integration_test.go

@@ -7,6 +7,8 @@ import (
 	"encoding/csv"
 	"fmt"
 	"iter"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -606,7 +608,7 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 			t.Helper()
 			var buf bytes.Buffer
 			iter := client.Single().Query(ctx, spanner.Statement{SQL: sql})
-			if err := writeCsvFromRowIter(&buf, iter, redact); err != nil {
+			if _, err := writeCsvFromRowIter(&buf, iter, redact); err != nil {
 				t.Fatalf("writeCsvFromRowIter(redact=%v): %v", redact, err)
 			}
 			return readCSV(t, buf.String())
@@ -645,4 +647,81 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 			}
 		})
 	})
+
+	t.Run("split plan output", func(t *testing.T) {
+		t.Setenv("SPANNER_EMULATOR_HOST", env.Emulator().URI())
+
+		t.Run("PROFILE csv publishes rows and errors without planNodes", func(t *testing.T) {
+			dir := t.TempDir()
+			rowsPath := filepath.Join(dir, "rows.csv")
+			planPath := filepath.Join(dir, "plan.json")
+			err := runMain(t, []string{
+				env.DatabaseID, "--project", env.ProjectID, "--instance", env.InstanceID,
+				"--sql", "SELECT SingerId, FirstName FROM Singers ORDER BY SingerId LIMIT 2",
+				"--query-mode", "PROFILE",
+				"--format", "experimental_csv",
+				"--output", rowsPath,
+				"--plan-output", planPath,
+			})
+			if err == nil || !strings.Contains(err.Error(), "no query plan") {
+				t.Fatalf("runMain() error = %v, want no query plan", err)
+			}
+			if !strings.Contains(err.Error(), "primary output written") {
+				t.Fatalf("error = %v, want primary published", err)
+			}
+			csvBytes, err := os.ReadFile(rowsPath)
+			if err != nil {
+				t.Fatal(err)
+			}
+			csvText := string(csvBytes)
+			if strings.Contains(csvText, "queryPlan") {
+				t.Fatalf("CSV contains queryPlan: %s", csvText)
+			}
+			recs, err := csv.NewReader(strings.NewReader(csvText)).ReadAll()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if len(recs) != 3 || recs[0][0] != "SingerId" {
+				t.Fatalf("csv records = %#v, want header + 2 rows", recs)
+			}
+			if _, err := os.Stat(planPath); !os.IsNotExist(err) {
+				t.Fatalf("plan file exists after missing-plan error: %v", err)
+			}
+		})
+
+		assertIteratorMetadata := func(t *testing.T, sql string, mode sppb.ExecuteSqlRequest_QueryMode, wantHeader []string, wantRows int) {
+			t.Helper()
+			var buf bytes.Buffer
+			result, err := writeCsvFromRowIter(&buf,
+				client.Single().QueryWithOptions(ctx, spanner.Statement{SQL: sql}, spanner.QueryOptions{Mode: mode.Enum()}),
+				false,
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if result == nil || result.Metadata == nil || result.Metadata.GetRowType() == nil {
+				t.Fatalf("missing metadata: %#v", result)
+			}
+			recs, err := csv.NewReader(strings.NewReader(buf.String())).ReadAll()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if len(recs) != 1+wantRows {
+				t.Fatalf("csv records = %#v, want header + %d rows", recs, wantRows)
+			}
+			if !cmp.Equal(recs[0], wantHeader) {
+				t.Fatalf("header = %v, want %v", recs[0], wantHeader)
+			}
+		}
+
+		t.Run("zero_row PROFILE still has metadata", func(t *testing.T) {
+			assertIteratorMetadata(t, "SELECT SingerId FROM Singers WHERE SingerId = -1",
+				sppb.ExecuteSqlRequest_PROFILE, []string{"SingerId"}, 0)
+		})
+
+		t.Run("PLAN still has metadata", func(t *testing.T) {
+			assertIteratorMetadata(t, "SELECT SingerId FROM Singers",
+				sppb.ExecuteSqlRequest_PLAN, []string{"SingerId"}, 0)
+		})
+	})
 }

jqresult/lazy.go

@@ -27,11 +27,29 @@ type Lazy struct {
 
 	drained  bool
 	drainErr error
+
+	omitQueryPlan bool
+}
+
+// LazyOption configures NewLazy.
+type LazyOption func(*Lazy)
+
+// WithOmitQueryPlan omits queryPlan from the .stats object presented to jq.
+// The underlying iterator stats still include the plan for Drain/Result.
+func WithOmitQueryPlan() LazyOption {
+	return func(l *Lazy) {
+		l.omitQueryPlan = true
+	}
 }
 
 // NewLazy builds a lazy jq input. rowIter must not have been read yet; Lazy takes ownership and Stop()s it.
-func NewLazy(rowIter *spanner.RowIterator, redact bool) *Lazy {
+func NewLazy(rowIter *spanner.RowIterator, redact bool, opts ...LazyOption) *Lazy {
 	l := &Lazy{redact: redact}
+	for _, opt := range opts {
+		if opt != nil {
+			opt(l)
+		}
+	}
 	l.rows = NewRowIter(rowIter, redact, RowToJSON)
 	l.rows.ioMu = &l.ioMu
 	return l
@@ -89,10 +107,79 @@ func (l *Lazy) drain() error {
 }
 
 func (l *Lazy) statsMapFromResult(result spaniter.RowIteratorResult) (map[string]any, error) {
+	var (
+		stats map[string]any
+		err   error
+	)
 	if l.encodeStats != nil {
-		return l.encodeStats(result.Stats)
+		stats, err = l.encodeStats(result.Stats)
+	} else {
+		stats, err = StatsMapFromResult(result)
+	}
+	if err != nil || stats == nil || !l.omitQueryPlan {
+		return stats, err
+	}
+	out := make(map[string]any, len(stats))
+	for k, v := range stats {
+		if k == "queryPlan" {
+			continue
+		}
+		out[k] = v
+	}
+	return out, nil
+}
+
+// Drain consumes any remaining rows so final stats are available.
+// If .stats was already read, this is a no-op.
+// Newly consumed rows are discarded and not retained on Lazy.
+func (l *Lazy) Drain() error {
+	l.mu.Lock()
+	if l.drained {
+		err := l.drainErr
+		l.mu.Unlock()
+		return err
+	}
+	alreadyStreamed := l.rowsStreamDone
+	l.mu.Unlock()
+
+	l.ioMu.Lock()
+	l.mu.Lock()
+	if l.drained {
+		err := l.drainErr
+		l.mu.Unlock()
+		l.ioMu.Unlock()
+		return err
+	}
+	l.mu.Unlock()
+
+	var drainErr error
+	if !alreadyStreamed {
+		drainErr = l.rows.discardRemainingUnlocked()
+	}
+	var stats map[string]any
+	if drainErr == nil {
+		stats, drainErr = l.statsMapFromResult(l.rows.Result())
+	}
+	l.ioMu.Unlock()
+
+	l.mu.Lock()
+	l.stats = stats
+	l.drainErr = drainErr
+	l.drained = true
+	l.rowsStreamDone = true
+	l.mu.Unlock()
+
+	l.rows.Stop()
+	return drainErr
+}
+
+// Result returns iterator metadata and stats captured while rows were consumed.
+// Call Drain first when jq may have stopped before the last PartialResultSet.
+func (l *Lazy) Result() spaniter.RowIteratorResult {
+	if l.rows == nil {
+		return spaniter.RowIteratorResult{}
 	}
-	return StatsMapFromResult(result)
+	return l.rows.Result()
 }
 
 func (l *Lazy) ensureMetadata() error {