diff --git a/README.md b/README.md index faebaee..052bd7d 100644 --- a/README.md +++ b/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`; 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`. diff --git a/flags_validation_test.go b/flags_validation_test.go index dabba5c..92ddbdb 100644 --- a/flags_validation_test.go +++ b/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 { diff --git a/integration_test.go b/integration_test.go index fbd88a9..8c6b756 100644 --- a/integration_test.go +++ b/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) + }) + }) } diff --git a/jqresult/lazy.go b/jqresult/lazy.go index 4cb2745..44b9f1e 100644 --- a/jqresult/lazy.go +++ b/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 { diff --git a/jqresult/lazy_stats_test.go b/jqresult/lazy_stats_test.go index 1f621b3..b56e578 100644 --- a/jqresult/lazy_stats_test.go +++ b/jqresult/lazy_stats_test.go @@ -52,6 +52,89 @@ func TestLazyRowsKeepsPositionAfterStatsDrain(t *testing.T) { } } +func TestLazyOmitQueryPlanLeavesQueryStats(t *testing.T) { + t.Parallel() + + l := newSyntheticLazy(t, 2) + l.omitQueryPlan = true + l.encodeStats = func(spaniter.Stats) (map[string]any, error) { + return map[string]any{ + "queryPlan": map[string]any{"planNodes": []any{map[string]any{"displayName": "Scan"}}}, + "queryStats": map[string]any{"elapsed_time": "1 msecs"}, + }, nil + } + defer l.Stop() + + code, err := Compile(".stats", InputLazy) + if err != nil { + t.Fatal(err) + } + iter := code.Run(l) + v, ok := iter.Next() + if !ok { + t.Fatal("no stats") + } + if err, isErr := v.(error); isErr { + t.Fatal(err) + } + stats, ok := v.(map[string]any) + if !ok { + t.Fatalf("stats type %T", v) + } + if _, ok := stats["queryPlan"]; ok { + t.Fatalf("queryPlan present: %#v", stats) + } + qs, ok := stats["queryStats"].(map[string]any) + if !ok || qs["elapsed_time"] != "1 msecs" { + t.Fatalf("queryStats = %#v", stats["queryStats"]) + } +} + +func TestLazyDrainDiscardsRemainingRows(t *testing.T) { + t.Parallel() + + l := newSyntheticLazy(t, 5) + defer l.Stop() + f, ok := l.rowsJQValue().(*lazyRowsField) + if !ok { + t.Fatal("rows view") + } + if _, ok := f.Next(); !ok { + t.Fatal("expected first row") + } + if err := l.Drain(); err != nil { + t.Fatal(err) + } + l.mu.Lock() + n := len(l.materializedRows) + l.mu.Unlock() + if n != 1 { + t.Fatalf("retained %d rows after Drain, want 1", n) + } +} + +func TestLazyDrainIsNoopAfterStats(t *testing.T) { + t.Parallel() + + l := newSyntheticLazy(t, 3) + defer l.Stop() + if _, err := l.statsMap(); err != nil { + t.Fatal(err) + } + l.mu.Lock() + before := len(l.materializedRows) + l.mu.Unlock() + if err := l.Drain(); err != nil { + t.Fatal(err) + } + l.mu.Lock() + after := len(l.materializedRows) + l.mu.Unlock() + if after != before { + t.Fatalf("Drain after stats changed retained rows %d -> %d", before, after) + } +} + func TestLazyStatsInterleavedFilterIDs(t *testing.T) { t.Parallel() diff --git a/jqresult/pipeline.go b/jqresult/pipeline.go index d1e65c3..f9210fb 100644 --- a/jqresult/pipeline.go +++ b/jqresult/pipeline.go @@ -12,7 +12,8 @@ import ( // For lazy mode, rowIter must be unread; cleanup releases the iterator state. // Lazy mode is intended for read-only queries; read-write callers should // materialize first and use eager mode. -func Execute(code *gojq.Code, mode InputMode, rowIter *spanner.RowIterator, rs *sppb.ResultSet, redactRows bool) (gojq.Iter, func(), error) { +// opts apply only to lazy mode (for example WithOmitQueryPlan). +func Execute(code *gojq.Code, mode InputMode, rowIter *spanner.RowIterator, rs *sppb.ResultSet, redactRows bool, opts ...LazyOption) (gojq.Iter, func(), error) { switch mode { case InputEager: if rs == nil { @@ -27,7 +28,7 @@ func Execute(code *gojq.Code, mode InputMode, rowIter *spanner.RowIterator, rs * if rowIter == nil { return nil, func() {}, fmt.Errorf("lazy mode requires an unread RowIterator") } - lazy := NewLazy(rowIter, redactRows) + lazy := NewLazy(rowIter, redactRows, opts...) return code.Run(lazy), lazy.Stop, nil default: return nil, func() {}, fmt.Errorf("unknown jq input mode: %s", mode) diff --git a/jqresult/rowiter.go b/jqresult/rowiter.go index 82d4880..d2741e6 100644 --- a/jqresult/rowiter.go +++ b/jqresult/rowiter.go @@ -149,6 +149,22 @@ func (r *RowIter) drainUnlocked() ([]any, error) { } } +// discardRemainingUnlocked consumes remaining iterator rows without retaining them. +func (r *RowIter) discardRemainingUnlocked() error { + if r.stopped { + return nil + } + for { + row, err := r.nextRow() + if row == nil && err == nil { + return nil + } + if err != nil { + return err + } + } +} + func (r *RowIter) Stop() { unlock := r.lockIO() defer unlock() diff --git a/main.go b/main.go index b7caa4f..7aa9d5c 100644 --- a/main.go +++ b/main.go @@ -56,7 +56,11 @@ type opts struct { DatabaseRole string `name:"database-role" help:"Database role to assume for all operations."` 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."` + Format string `name:"format" enum:"json,yaml,experimental_csv" default:"json" help:"Output format of the primary document."` + Output string `name:"output" short:"o" default:"-" help:"Destination of the primary document. Use - for stdout; /dev/stdout and /dev/stderr are mapped in-process."` + PlanOutput string `name:"plan-output" help:"Write the query-plan artifact here and strip stats.queryPlan from the primary document. Enables split mode."` + PlanFormat string `name:"plan-format" help:"Format of the plan artifact: json or yaml. Defaults to --format when that is json or yaml, otherwise json. Requires --plan-output."` + DiscardResults bool `name:"discard-results" help:"Do not write the primary document (plan-only). Requires --plan-output."` 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)"` JqFilter string `name:"filter" xor:"filter" help:"jq filter"` @@ -228,6 +232,9 @@ func queryModeForQuery(query string, enablePartitionedDML bool, tb spanner.Times } func validateExecutionOptions(o opts, mode queryMode) error { + if err := validatePlanOutputOptions(o); err != nil { + return err + } if o.TryPartitionQuery { if _, ok := mode.(single); !ok { if o.EnablePartitionedDML { @@ -462,6 +469,12 @@ func runCLI(clientOptions ...option.ClientOption) error { return err } + sinks, err := newOutputSinks(o) + if err != nil { + return err + } + defer sinks.Abort() + ctx, tp, err := enableTracing(ctx, o) if err != nil { return err @@ -504,47 +517,155 @@ func runCLI(clientOptions ...option.ClientOption) error { return err } - fmt.Println("success") - return nil + if sinks.primary != nil { + if _, err := fmt.Fprintln(sinks.primary, "success"); err != nil { + return err + } + } + sinks.MarkPrimaryComplete() + return sinks.Finish(nil) } + var workErr error if o.Format == "experimental_csv" { - return runAndWriteCsv(ctx, client, stmt, queryOpts, m, o.RedactRows) + workErr = runAndWriteCsv(ctx, client, stmt, queryOpts, m, o, sinks) + } else { + workErr = runJqOutput(ctx, client, stmt, queryOpts, m, o, jqMode, jqCode, sinks) + } + finishErr := sinks.Finish(workErr) + if finishErr != nil && workErr == nil && isCommittedMode(m) { + // The statement completed (a read-write transaction committed or a + // partitioned DML finished) and only file publication failed. + // Say so explicitly so nobody replays the DML to repair an output file. + return wrapCommittedOutputError(finishErr) + } + return finishErr +} + +// isCommittedMode reports whether a successful run of mode leaves a committed +// write behind, which changes how later output failures must be described. +func isCommittedMode(mode queryMode) bool { + switch mode.(type) { + case readWrite, partitionedDML: + return true + default: + return false } +} - return runJqOutput(ctx, client, stmt, queryOpts, m, o, jqMode, jqCode) +// materializeWithoutRows reports whether the eager path may drop row values +// while materializing: both --redact-rows and --discard-results never emit +// rows, so reading them into memory would only cost time and memory. +func materializeWithoutRows(o opts) bool { + return o.RedactRows || o.DiscardResults } -func runAndWriteCsv(ctx context.Context, client *spanner.Client, stmt spanner.Statement, opts spanner.QueryOptions, mode queryMode, redactRows bool) error { +func runAndWriteCsv(ctx context.Context, client *spanner.Client, stmt spanner.Statement, opts spanner.QueryOptions, mode queryMode, o opts, sinks *outputSinks) error { + encodeRowCount := dmlRowCountForMode(mode, opts) + statOpts := spaniterStatsOpts(mode, opts) + planFmt := effectivePlanFormat(o) + writePlanFromCSV := func(result *svwriter.RowIteratorResult) error { + if sinks.plan == nil { + return nil + } + stats, err := statsFromWriterResult(result, encodeRowCount) + if err != nil { + return err + } + var md *sppb.ResultSetMetadata + if result != nil { + md = result.Metadata + } + return writePlan(sinks.plan, planFmt, md, stats) + } + writePlanFromDrain := func(result *spaniter.RowIteratorResult) error { + if sinks.plan == nil { + return nil + } + if result == nil { + return errNoQueryPlan + } + stats, err := result.StatsProto() + if err != nil { + return err + } + return writePlan(sinks.plan, planFmt, result.Metadata, stats) + } + switch mode := mode.(type) { case readWrite: var buf bytes.Buffer + var csvResult *svwriter.RowIteratorResult + var drainResult *spaniter.RowIteratorResult _, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *spanner.ReadWriteTransaction) error { buf.Reset() - return writeCsvFromRowIter(&buf, tx.QueryWithOptions(ctx, stmt, opts), redactRows) + rowIter := tx.QueryWithOptions(ctx, stmt, opts) + if o.DiscardResults { + var err error + drainResult, err = spaniter.DrainRowIterator(rowIter, statOpts...) + return err + } + var err error + csvResult, err = writeCsvFromRowIter(&buf, rowIter, o.RedactRows) + return err }) if err != nil { return err } - _, err = io.Copy(os.Stdout, &buf) - return err + if sinks.primary != nil { + if _, err := io.Copy(sinks.primary, &buf); err != nil { + return wrapCommittedOutputError(err) + } + } + sinks.MarkPrimaryComplete() + var planErr error + if o.DiscardResults { + planErr = writePlanFromDrain(drainResult) + } else { + planErr = writePlanFromCSV(csvResult) + } + if planErr != nil { + return wrapCommittedOutputError(planErr) + } + return nil case single: - return writeCsvFromRowIter( - os.Stdout, - client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, stmt, opts), - redactRows, - ) + rowIter := client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, stmt, opts) + if o.DiscardResults { + result, err := spaniter.DrainRowIterator(rowIter, statOpts...) + if err != nil { + return err + } + sinks.MarkPrimaryComplete() + return writePlanFromDrain(result) + } + writer := sinks.primary + if writer == nil { + writer = io.Discard + } + result, err := writeCsvFromRowIter(writer, rowIter, o.RedactRows) + if err != nil { + return err + } + sinks.MarkPrimaryComplete() + return writePlanFromCSV(result) case partitionedDML: count, err := client.PartitionedUpdateWithOptions(ctx, stmt, opts) if err != nil { return err } - return writeCsvFromResultSet(os.Stdout, &sppb.ResultSet{ + rs := &sppb.ResultSet{ Metadata: &sppb.ResultSetMetadata{RowType: &sppb.StructType{}}, Stats: &sppb.ResultSetStats{ RowCount: &sppb.ResultSetStats_RowCountLowerBound{RowCountLowerBound: count}, }, - }) + } + if sinks.primary != nil { + if err := writeCsvFromResultSet(sinks.primary, rs); err != nil { + return wrapCommittedOutputError(err) + } + } + sinks.MarkPrimaryComplete() + return nil default: panic(fmt.Sprintf("unknown mode: %T", mode)) } @@ -561,17 +682,16 @@ func (csvRedactRowIteratorWriter) WriteRow(*spanner.Row) error { return nil } // writeCsvFromRowIter streams query rows to CSV without materializing a ResultSet. // Pass the query iterator directly to WriteRowIterator (it owns Stop); do not defer Stop at the call site. -func writeCsvFromRowIter(writer io.Writer, rowIter *spanner.RowIterator, redactRows bool) error { +func writeCsvFromRowIter(writer io.Writer, rowIter *spanner.RowIterator, redactRows bool) (*svwriter.RowIteratorResult, error) { csvWriter, err := svwriter.NewCSVWriter(writer) if err != nil { - return err + return nil, err } iterWriter := svwriter.RowIteratorWriter(csvWriter) if redactRows { iterWriter = csvRedactRowIteratorWriter{csvWriter} } - _, err = svwriter.WriteRowIterator(rowIter, iterWriter) - return err + return svwriter.WriteRowIterator(rowIter, iterWriter) } func prepareCsvRowType(csvWriter *svwriter.DelimitedWriter, metadata *sppb.ResultSetMetadata) error { @@ -657,38 +777,88 @@ func runJqOutput( o opts, jqMode jqresult.InputMode, jqCode *gojq.Code, + sinks *outputSinks, ) error { + committed := isCommittedMode(mode) + wrap := func(err error) error { + if err != nil && committed { + return wrapCommittedOutputError(err) + } + return err + } useEager := jqMode == jqresult.InputEager // Read-write DML always materializes the full result set before jq runs. if _, ok := mode.(readWrite); ok { useEager = true } + planFmt := effectivePlanFormat(o) if useEager { - rs, err := runInNewTransaction(ctx, client, stmt, opts, mode, o.RedactRows) + rs, err := runInNewTransaction(ctx, client, stmt, opts, mode, materializeWithoutRows(o)) if err != nil { return err } - enc, err := newEncoder(os.Stdout, o.Format, o.CompactOutput, o.JqRawOutput) - if err != nil { - return err + var planStats *sppb.ResultSetStats + var metadata *sppb.ResultSetMetadata + if sinks.hasPlan { + planStats, metadata = stripQueryPlanForPrimary(rs) + } else if rs != nil { + metadata = rs.Metadata + planStats = rs.Stats } - defer func() { _ = closeEncoder(enc) }() - iter, cleanup, err := jqresult.Execute(jqCode, jqresult.InputEager, nil, rs, o.RedactRows) - if err != nil { - return err + if sinks.primary != nil { + enc, err := newEncoder(sinks.primary, o.Format, o.CompactOutput, o.JqRawOutput) + if err != nil { + return wrap(err) + } + iter, cleanup, err := jqresult.Execute(jqCode, jqresult.InputEager, nil, rs, o.RedactRows) + if err != nil { + _ = closeEncoder(enc) + return wrap(err) + } + printErr := jqresult.Print(enc, iter) + cleanup() + closeErr := closeEncoder(enc) + if printErr != nil { + return wrap(printErr) + } + if closeErr != nil { + return wrap(closeErr) + } } - defer cleanup() - return jqresult.Print(enc, iter) + sinks.MarkPrimaryComplete() + if sinks.plan != nil { + return wrap(writePlan(sinks.plan, planFmt, metadata, planStats)) + } + return nil } switch mode := mode.(type) { case single: - enc, err := newEncoder(os.Stdout, o.Format, o.CompactOutput, o.JqRawOutput) + rowIter := client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, stmt, opts) + if o.DiscardResults { + result, err := spaniter.DrainRowIterator(rowIter, spaniterStatsOpts(mode, opts)...) + if err != nil { + return err + } + sinks.MarkPrimaryComplete() + if sinks.plan == nil { + return nil + } + stats, err := result.StatsProto() + if err != nil { + return err + } + return writePlan(sinks.plan, planFmt, result.Metadata, stats) + } + writer := sinks.primary + if writer == nil { + writer = io.Discard + } + enc, err := newEncoder(writer, o.Format, o.CompactOutput, o.JqRawOutput) if err != nil { return err } - rowIter := client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, stmt, opts) - return runJqOnRowIter(rowIter, o.RedactRows, jqCode, enc) + return runJqOnRowIter(rowIter, o.RedactRows, jqCode, enc, sinks, planFmt) case partitionedDML: return fmt.Errorf("--jq-input-mode=lazy is not supported for partitioned DML") default: @@ -701,14 +871,36 @@ func runJqOnRowIter( redactRows bool, jqCode *gojq.Code, enc encoder, + sinks *outputSinks, + planFmt string, ) error { - defer func() { _ = closeEncoder(enc) }() - iter, cleanup, err := jqresult.Execute(jqCode, jqresult.InputLazy, rowIter, nil, redactRows) + var lazyOpts []jqresult.LazyOption + if sinks.hasPlan { + lazyOpts = append(lazyOpts, jqresult.WithOmitQueryPlan()) + } + lazy := jqresult.NewLazy(rowIter, redactRows, lazyOpts...) + defer lazy.Stop() + printErr := jqresult.Print(enc, jqCode.Run(lazy)) + closeErr := closeEncoder(enc) + if printErr != nil { + return printErr + } + if closeErr != nil { + return closeErr + } + sinks.MarkPrimaryComplete() + if sinks.plan == nil { + return nil + } + if err := lazy.Drain(); err != nil { + return err + } + result := lazy.Result() + stats, err := result.StatsProto() if err != nil { return err } - defer cleanup() - return jqresult.Print(enc, iter) + return writePlan(sinks.plan, planFmt, result.Metadata, stats) } func newEncoder(writer io.Writer, format string, compactOutput bool, rawOutput bool) (encoder, error) { diff --git a/output.go b/output.go new file mode 100644 index 0000000..5a56a29 --- /dev/null +++ b/output.go @@ -0,0 +1,482 @@ +package main + +import ( + "errors" + "fmt" + "io" + "os" + "path/filepath" + + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/jqresult" + svwriter "github.com/apstndb/spanvalue/writer" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/structpb" +) + +const ( + destStdoutDash = "-" + destDevStdout = "/dev/stdout" + destDevStderr = "/dev/stderr" + + planModesHelp = "PLAN, PROFILE, or WITH_PLAN_AND_STATS" +) + +var errNoQueryPlan = errors.New("query returned no query plan (nil or empty planNodes); not retrying the statement") + +type destKind int + +const ( + destKindStdout destKind = iota + destKindStderr + destKindFile +) + +type resolvedDest struct { + kind destKind + raw string + abs string +} + +func defaultPrimaryOutput(raw string) string { + if raw == "" { + return destStdoutDash + } + return raw +} + +func resolveDestination(raw string) resolvedDest { + switch raw { + case destStdoutDash, destDevStdout: + return resolvedDest{kind: destKindStdout, raw: raw} + case destDevStderr: + return resolvedDest{kind: destKindStderr, raw: raw} + default: + abs, err := filepath.Abs(filepath.Clean(raw)) + if err != nil { + abs = filepath.Clean(raw) + } + return resolvedDest{kind: destKindFile, raw: raw, abs: abs} + } +} + +func destLabel(d resolvedDest) string { + switch d.kind { + case destKindStdout: + return "stdout" + case destKindStderr: + return "stderr" + default: + if d.abs != "" { + return d.abs + } + return d.raw + } +} + +func isPlanProducingQueryMode(mode string) bool { + switch mode { + case "PLAN", "PROFILE", "WITH_PLAN_AND_STATS": + return true + default: + return false + } +} + +func effectivePlanFormat(o opts) string { + if o.PlanFormat != "" { + return o.PlanFormat + } + switch o.Format { + case "json", "yaml": + return o.Format + default: + return "json" + } +} + +func validatePlanOutputOptions(o opts) error { + hasPlan := o.PlanOutput != "" + if o.PlanFormat != "" && !hasPlan { + return fmt.Errorf("--plan-format requires --plan-output") + } + if o.DiscardResults && !hasPlan { + return fmt.Errorf("--discard-results requires --plan-output") + } + if o.PlanFormat != "" && o.PlanFormat != "json" && o.PlanFormat != "yaml" { + return fmt.Errorf("--plan-format must be json or yaml") + } + if !hasPlan { + return nil + } + mode := o.QueryMode + if mode == "" { + mode = "NORMAL" + } + if !isPlanProducingQueryMode(mode) { + return fmt.Errorf("--plan-output requires --query-mode=%s", planModesHelp) + } + if o.TryPartitionQuery { + return fmt.Errorf("--plan-output cannot be combined with --try-partition-query") + } + if o.EnablePartitionedDML { + return fmt.Errorf("--plan-output cannot be combined with --enable-partitioned-dml") + } + return nil +} + +func validateDestinations(o opts) error { + primaryRaw := defaultPrimaryOutput(o.Output) + primary := resolveDestination(primaryRaw) + var plan *resolvedDest + if o.PlanOutput != "" { + d := resolveDestination(o.PlanOutput) + plan = &d + } + + if plan != nil && !o.DiscardResults { + if primary.kind == destKindStdout && plan.kind == destKindStdout { + return fmt.Errorf("in split mode --output and --plan-output cannot both write to stdout") + } + if primary.kind == destKindStderr && plan.kind == destKindStderr { + return fmt.Errorf("in split mode --output and --plan-output cannot both write to stderr") + } + if primary.kind == destKindFile && plan.kind == destKindFile { + same, err := sameOutputFile(primary.raw, plan.raw) + if err != nil { + return err + } + if same { + return fmt.Errorf("--output and --plan-output cannot target the same file") + } + } + } + + inputs := []struct { + flag, path string + }{ + {"--sql-file", o.SqlFile}, + {"--param-file", o.ParamFile}, + {"--filter-file", o.JqFromFile}, + } + + checkAlias := func(flagName, destRaw string, d resolvedDest) error { + if d.kind != destKindFile { + return nil + } + for _, in := range inputs { + if in.path == "" { + continue + } + same, err := sameOutputFile(destRaw, in.path) + if err != nil { + return err + } + if same { + return fmt.Errorf("%s cannot alias %s", flagName, in.flag) + } + } + return nil + } + + if !o.DiscardResults { + if err := checkAlias("--output", primaryRaw, primary); err != nil { + return err + } + } + if plan != nil { + if err := checkAlias("--plan-output", o.PlanOutput, *plan); err != nil { + return err + } + } + return nil +} + +func fileIdentity(path string) (abs string, info os.FileInfo, exists bool, err error) { + abs, err = filepath.Abs(filepath.Clean(path)) + if err != nil { + return "", nil, false, err + } + info, statErr := os.Stat(path) + if statErr != nil { + if errors.Is(statErr, os.ErrNotExist) { + return abs, nil, false, nil + } + return abs, nil, false, statErr + } + if eval, evalErr := filepath.EvalSymlinks(path); evalErr == nil { + if evalAbs, absErr := filepath.Abs(eval); absErr == nil { + abs = evalAbs + } + } + return abs, info, true, nil +} + +func sameOutputFile(a, b string) (bool, error) { + absA, infoA, existsA, err := fileIdentity(a) + if err != nil { + return false, err + } + absB, infoB, existsB, err := fileIdentity(b) + if err != nil { + return false, err + } + if absA == absB { + return true, nil + } + if existsA && existsB { + return os.SameFile(infoA, infoB), nil + } + return false, nil +} + +// fileSink is a regular-file destination written via a sibling temp file. +type fileSink struct { + file *os.File + final string +} + +type outputSinks struct { + primary io.Writer + plan io.Writer + + primaryDest resolvedDest + planDest resolvedDest + hasPlan bool + discard bool + + primaryFile *fileSink + planFile *fileSink + + primaryReady bool + done bool +} + +func newOutputSinks(o opts) (*outputSinks, error) { + if err := validateDestinations(o); err != nil { + return nil, err + } + s := &outputSinks{discard: o.DiscardResults} + if o.DiscardResults { + s.primaryReady = true + } else { + d := resolveDestination(defaultPrimaryOutput(o.Output)) + w, file, err := openDestination(d) + if err != nil { + return nil, fmt.Errorf("--output: %w", err) + } + s.primary = w + s.primaryDest = d + s.primaryFile = file + } + if o.PlanOutput != "" { + d := resolveDestination(o.PlanOutput) + w, file, err := openDestination(d) + if err != nil { + s.Abort() + return nil, fmt.Errorf("--plan-output: %w", err) + } + s.plan = w + s.planDest = d + s.planFile = file + s.hasPlan = true + } + return s, nil +} + +func openDestination(d resolvedDest) (io.Writer, *fileSink, error) { + switch d.kind { + case destKindStdout: + return os.Stdout, nil, nil + case destKindStderr: + return os.Stderr, nil, nil + case destKindFile: + dir := filepath.Dir(d.abs) + tmp, err := os.CreateTemp(dir, ".execspansql-*.tmp") + if err != nil { + return nil, nil, err + } + if err := tmp.Chmod(0o600); err != nil { + name := tmp.Name() + _ = tmp.Close() + _ = os.Remove(name) + return nil, nil, err + } + return tmp, &fileSink{file: tmp, final: d.abs}, nil + default: + return nil, nil, fmt.Errorf("unknown destination kind") + } +} + +func (s *outputSinks) MarkPrimaryComplete() { + if s == nil { + return + } + s.primaryReady = true +} + +func (s *outputSinks) primaryLabel() string { + if s == nil || s.discard { + return "" + } + return destLabel(s.primaryDest) +} + +func abortFileSink(fs *fileSink) { + if fs == nil || fs.file == nil { + return + } + name := fs.file.Name() + _ = fs.file.Close() + _ = os.Remove(name) + fs.file = nil +} + +func publishFileSink(fs *fileSink) error { + if fs == nil || fs.file == nil { + return nil + } + name := fs.file.Name() + if err := fs.file.Close(); err != nil { + _ = os.Remove(name) + fs.file = nil + return err + } + fs.file = nil + if err := os.Rename(name, fs.final); err != nil { + _ = os.Remove(name) + return err + } + return nil +} + +func (s *outputSinks) Abort() { + if s == nil || s.done { + return + } + s.done = true + abortFileSink(s.primaryFile) + abortFileSink(s.planFile) +} + +func (s *outputSinks) Finish(workErr error) error { + if s == nil { + return workErr + } + if workErr != nil && !s.primaryReady { + s.Abort() + return workErr + } + if workErr != nil { + pubErr := publishFileSink(s.primaryFile) + abortFileSink(s.planFile) + s.done = true + if pubErr != nil { + return fmt.Errorf("%v; also failed to publish primary output: %w", workErr, pubErr) + } + if label := s.primaryLabel(); label != "" { + return fmt.Errorf("primary output written to %s; %w", label, workErr) + } + return workErr + } + if err := publishFileSink(s.primaryFile); err != nil { + abortFileSink(s.planFile) + s.done = true + return err + } + if err := publishFileSink(s.planFile); err != nil { + s.done = true + if label := s.primaryLabel(); label != "" { + return fmt.Errorf("primary output written to %s; failed to publish plan output: %w", label, err) + } + return fmt.Errorf("failed to publish plan output: %w", err) + } + s.done = true + return nil +} + +func wrapCommittedOutputError(err error) error { + if err == nil { + return nil + } + return fmt.Errorf("output failed after the statement was committed; this is not a rollback and the SQL is not replayed: %w", err) +} + +func hasUsableQueryPlan(stats *sppb.ResultSetStats) bool { + if stats == nil || stats.GetQueryPlan() == nil { + return false + } + return len(stats.GetQueryPlan().GetPlanNodes()) > 0 +} + +// stripQueryPlanForPrimary clones stats for the plan artifact and clears QueryPlan +// on a separate clone assigned to rs. The original stats message is not mutated. +func stripQueryPlanForPrimary(rs *sppb.ResultSet) (planStats *sppb.ResultSetStats, metadata *sppb.ResultSetMetadata) { + if rs == nil { + return nil, nil + } + metadata = rs.Metadata + if rs.Stats == nil { + return nil, metadata + } + planStats, ok := proto.Clone(rs.Stats).(*sppb.ResultSetStats) + if !ok || planStats == nil { + return nil, metadata + } + primaryStats, ok := proto.Clone(rs.Stats).(*sppb.ResultSetStats) + if !ok || primaryStats == nil { + return planStats, metadata + } + primaryStats.QueryPlan = nil + rs.Stats = primaryStats + return planStats, metadata +} + +func writePlan(w io.Writer, format string, metadata *sppb.ResultSetMetadata, stats *sppb.ResultSetStats) error { + if w == nil { + return nil + } + if !hasUsableQueryPlan(stats) { + return errNoQueryPlan + } + envelope := &sppb.ResultSet{ + Metadata: metadata, + Stats: stats, + } + m, err := jqresult.ProtoToMap(envelope) + if err != nil { + return err + } + enc, err := newEncoder(w, format, false, false) + if err != nil { + return err + } + if err := enc.Encode(m); err != nil { + _ = closeEncoder(enc) + return err + } + return closeEncoder(enc) +} + +func statsFromWriterResult(r *svwriter.RowIteratorResult, encodeRowCount bool) (*sppb.ResultSetStats, error) { + if r == nil { + return nil, nil + } + stats := &sppb.ResultSetStats{ + QueryPlan: r.Stats.QueryPlan, + } + if r.Stats.QueryStats != nil { + qs, err := structpb.NewStruct(r.Stats.QueryStats) + if err != nil { + return nil, fmt.Errorf("encode query stats: %w", err) + } + stats.QueryStats = qs + } + if encodeRowCount { + stats.RowCount = &sppb.ResultSetStats_RowCountExact{RowCountExact: r.Stats.RowCount} + } + if stats.QueryPlan == nil && stats.QueryStats == nil && stats.RowCount == nil { + return nil, nil + } + return stats, nil +} diff --git a/output_test.go b/output_test.go new file mode 100644 index 0000000..18dc567 --- /dev/null +++ b/output_test.go @@ -0,0 +1,571 @@ +package main + +import ( + "bytes" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/jqresult" + "google.golang.org/protobuf/types/known/structpb" +) + +func TestResolveDestinationStdoutSpellings(t *testing.T) { + t.Parallel() + + tests := []struct { + raw string + kind destKind + }{ + {raw: "-", kind: destKindStdout}, + {raw: "/dev/stdout", kind: destKindStdout}, + {raw: "/dev/stderr", kind: destKindStderr}, + {raw: "plan.json", kind: destKindFile}, + } + for _, tt := range tests { + t.Run(tt.raw, func(t *testing.T) { + got := resolveDestination(tt.raw) + if got.kind != tt.kind { + t.Fatalf("resolveDestination(%q).kind = %v, want %v", tt.raw, got.kind, tt.kind) + } + }) + } +} + +func TestEffectivePlanFormat(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + o opts + want string + }{ + {name: "explicit_yaml", o: opts{PlanFormat: "yaml", Format: "json"}, want: "yaml"}, + {name: "follows_json", o: opts{Format: "json"}, want: "json"}, + {name: "follows_yaml", o: opts{Format: "yaml"}, want: "yaml"}, + {name: "csv_defaults_to_json", o: opts{Format: "experimental_csv"}, want: "json"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := effectivePlanFormat(tt.o); got != tt.want { + t.Fatalf("effectivePlanFormat() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestValidateDestinationsStdoutCollision(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + o opts + err string + }{ + { + name: "dash_and_dash", + o: opts{Output: "-", PlanOutput: "-"}, + err: "cannot both write to stdout", + }, + { + name: "dash_and_dev_stdout", + o: opts{Output: "-", PlanOutput: "/dev/stdout"}, + err: "cannot both write to stdout", + }, + { + name: "dev_stdout_and_dash", + o: opts{Output: "/dev/stdout", PlanOutput: "-"}, + err: "cannot both write to stdout", + }, + { + name: "both_stderr", + o: opts{Output: "/dev/stderr", PlanOutput: "/dev/stderr"}, + err: "cannot both write to stderr", + }, + { + name: "stdout_and_stderr_ok", + o: opts{Output: "-", PlanOutput: "/dev/stderr"}, + }, + { + name: "discard_allows_plan_on_stdout", + o: opts{Output: "-", PlanOutput: "-", DiscardResults: true}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateDestinations(tt.o) + if tt.err == "" { + if err != nil { + t.Fatalf("validateDestinations() error = %v, want nil", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.err) { + t.Fatalf("validateDestinations() error = %v, want %q", err, tt.err) + } + }) + } +} + +func TestValidateDestinationsSameFile(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + path := filepath.Join(dir, "out.json") + if err := os.WriteFile(path, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + + err := validateDestinations(opts{Output: path, PlanOutput: path}) + if err == nil || !strings.Contains(err.Error(), "same file") { + t.Fatalf("same path: error = %v, want same file", err) + } + + rel := filepath.Join(dir, ".", "out.json") + err = validateDestinations(opts{Output: path, PlanOutput: rel}) + if err == nil || !strings.Contains(err.Error(), "same file") { + t.Fatalf("cleaned path: error = %v, want same file", err) + } +} + +func TestValidateDestinationsSymlinkAndHardLink(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + target := filepath.Join(dir, "plan.json") + if err := os.WriteFile(target, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + + link := filepath.Join(dir, "plan-link.json") + if err := os.Symlink(target, link); err != nil { + t.Fatal(err) + } + err := validateDestinations(opts{Output: target, PlanOutput: link}) + if err == nil || !strings.Contains(err.Error(), "same file") { + t.Fatalf("symlink: error = %v, want same file", err) + } + + hard := filepath.Join(dir, "plan-hard.json") + if err := os.Link(target, hard); err != nil { + t.Skipf("hard link not supported: %v", err) + } + err = validateDestinations(opts{Output: target, PlanOutput: hard}) + if err == nil || !strings.Contains(err.Error(), "same file") { + t.Fatalf("hard link: error = %v, want same file", err) + } +} + +func TestValidateDestinationsInputAlias(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + sqlFile := filepath.Join(dir, "query.sql") + paramFile := filepath.Join(dir, "params.yaml") + filterFile := filepath.Join(dir, "filter.jq") + for _, p := range []string{sqlFile, paramFile, filterFile} { + if err := os.WriteFile(p, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + } + + tests := []struct { + name string + o opts + err string + }{ + { + name: "output_aliases_sql_file", + o: opts{Output: sqlFile, SqlFile: sqlFile}, + err: "--output cannot alias --sql-file", + }, + { + name: "plan_aliases_param_file", + o: opts{Output: filepath.Join(dir, "rows.json"), PlanOutput: paramFile, ParamFile: paramFile}, + err: "--plan-output cannot alias --param-file", + }, + { + name: "output_aliases_filter_file", + o: opts{Output: filterFile, JqFromFile: filterFile}, + err: "--output cannot alias --filter-file", + }, + { + name: "discard_skips_primary_alias", + o: opts{ + Output: sqlFile, + PlanOutput: filepath.Join(dir, "plan.json"), + SqlFile: sqlFile, + DiscardResults: true, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateDestinations(tt.o) + if tt.err == "" { + if err != nil { + t.Fatalf("validateDestinations() error = %v, want nil", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.err) { + t.Fatalf("validateDestinations() error = %v, want %q", err, tt.err) + } + }) + } +} + +func TestStripQueryPlanForPrimaryDoesNotMutateOriginal(t *testing.T) { + t.Parallel() + + rs := profileResultSetForSplitTest() + orig := rs.Stats + planStats, md := stripQueryPlanForPrimary(rs) + if orig.GetQueryPlan() == nil || len(orig.GetQueryPlan().GetPlanNodes()) == 0 { + t.Fatal("original stats.QueryPlan was mutated") + } + if rs.Stats == orig { + t.Fatal("primary stats still shares the original stats pointer") + } + if rs.Stats.GetQueryPlan() != nil { + t.Fatal("primary stats still has QueryPlan") + } + if planStats.GetQueryPlan() == nil || planStats.GetQueryPlan().GetPlanNodes()[0].GetDisplayName() != "Scan" { + t.Fatal("plan clone missing QueryPlan") + } + if md == nil || md.GetRowType() == nil { + t.Fatal("metadata missing") + } + + m, err := jqresult.ResultSetMap(rs) + if err != nil { + t.Fatal(err) + } + stats, _ := m["stats"].(map[string]any) + if stats == nil { + t.Fatal("primary stats missing") + } + if _, ok := stats["queryPlan"]; ok { + t.Fatalf("primary stats still has queryPlan: %#v", stats) + } + if _, ok := stats["queryStats"]; !ok { + t.Fatalf("primary stats missing queryStats: %#v", stats) + } +} + +func TestWritePlanEnvelopeOmitsRows(t *testing.T) { + t.Parallel() + + rs := profileResultSetForSplitTest() + planStats, md := stripQueryPlanForPrimary(rs) + + var buf bytes.Buffer + if err := writePlan(&buf, "json", md, planStats); err != nil { + t.Fatal(err) + } + m, err := jqresult.ProtoToMap(&sppb.ResultSet{Metadata: md, Stats: planStats}) + if err != nil { + t.Fatal(err) + } + if _, ok := m["rows"]; ok { + t.Fatalf("plan envelope has rows: %#v", m) + } + stats, _ := m["stats"].(map[string]any) + if _, ok := stats["queryPlan"]; !ok { + t.Fatal("plan envelope missing queryPlan") + } + + var yamlBuf bytes.Buffer + if err := writePlan(&yamlBuf, "yaml", md, planStats); err != nil { + t.Fatal(err) + } + if !strings.Contains(yamlBuf.String(), "queryPlan") { + t.Fatalf("yaml plan missing queryPlan: %s", yamlBuf.String()) + } +} + +func TestWritePlanRejectsEmptyPlan(t *testing.T) { + t.Parallel() + + err := writePlan(ioDiscardWriter{}, "json", nil, &sppb.ResultSetStats{ + QueryPlan: &sppb.QueryPlan{}, + }) + if !errors.Is(err, errNoQueryPlan) { + t.Fatalf("error = %v, want errNoQueryPlan", err) + } +} + +func TestOutputSinksPublishAndAbort(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + primary := filepath.Join(dir, "rows.json") + if err := os.WriteFile(primary, []byte("old-primary"), 0o644); err != nil { + t.Fatal(err) + } + + t.Run("abort_leaves_existing", func(t *testing.T) { + s, err := newOutputSinks(opts{Output: primary}) + if err != nil { + t.Fatal(err) + } + if _, err := s.primary.Write([]byte("new-primary")); err != nil { + t.Fatal(err) + } + s.Abort() + got, err := os.ReadFile(primary) + if err != nil { + t.Fatal(err) + } + if string(got) != "old-primary" { + t.Fatalf("after abort: %q, want old-primary", got) + } + }) + + t.Run("finish_error_before_primary_ready_aborts", func(t *testing.T) { + s, err := newOutputSinks(opts{Output: primary}) + if err != nil { + t.Fatal(err) + } + if _, err := s.primary.Write([]byte("partial")); err != nil { + t.Fatal(err) + } + err = s.Finish(errors.New("query failed")) + if err == nil || !strings.Contains(err.Error(), "query failed") { + t.Fatalf("Finish() error = %v", err) + } + got, err := os.ReadFile(primary) + if err != nil { + t.Fatal(err) + } + if string(got) != "old-primary" { + t.Fatalf("after failed query: %q, want old-primary", got) + } + }) + + t.Run("publish_overwrites", func(t *testing.T) { + s, err := newOutputSinks(opts{Output: primary}) + if err != nil { + t.Fatal(err) + } + if _, err := s.primary.Write([]byte("new-primary")); err != nil { + t.Fatal(err) + } + s.MarkPrimaryComplete() + if err := s.Finish(nil); err != nil { + t.Fatal(err) + } + got, err := os.ReadFile(primary) + if err != nil { + t.Fatal(err) + } + if string(got) != "new-primary" { + t.Fatalf("after publish: %q, want new-primary", got) + } + }) + + t.Run("plan_error_publishes_primary", func(t *testing.T) { + plan := filepath.Join(dir, "plan.json") + if err := os.WriteFile(plan, []byte("old-plan"), 0o644); err != nil { + t.Fatal(err) + } + s, err := newOutputSinks(opts{Output: primary, PlanOutput: plan}) + if err != nil { + t.Fatal(err) + } + if _, err := s.primary.Write([]byte("rows-ok")); err != nil { + t.Fatal(err) + } + s.MarkPrimaryComplete() + err = s.Finish(errNoQueryPlan) + if err == nil || !strings.Contains(err.Error(), "primary output written") { + t.Fatalf("Finish() error = %v, want primary written", err) + } + got, err := os.ReadFile(primary) + if err != nil { + t.Fatal(err) + } + if string(got) != "rows-ok" { + t.Fatalf("primary = %q, want rows-ok", got) + } + gotPlan, err := os.ReadFile(plan) + if err != nil { + t.Fatal(err) + } + if string(gotPlan) != "old-plan" { + t.Fatalf("plan = %q, want old-plan", gotPlan) + } + }) +} + +func TestDiscardResultsProducesNoPrimaryBytes(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + primary := filepath.Join(dir, "rows.csv") + plan := filepath.Join(dir, "plan.json") + s, err := newOutputSinks(opts{ + Output: primary, + PlanOutput: plan, + DiscardResults: true, + QueryMode: "PROFILE", + }) + if err != nil { + t.Fatal(err) + } + if s.primary != nil { + t.Fatal("primary writer should be nil when discarding results") + } + if _, err := s.plan.Write([]byte("plan-bytes")); err != nil { + t.Fatal(err) + } + if err := s.Finish(nil); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(primary); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("primary file exists: %v", err) + } + got, err := os.ReadFile(plan) + if err != nil { + t.Fatal(err) + } + if string(got) != "plan-bytes" { + t.Fatalf("plan = %q", got) + } +} + +func TestIsCommittedMode(t *testing.T) { + t.Parallel() + + if isCommittedMode(single{}) { + t.Fatal("single-use read must not count as committed") + } + if !isCommittedMode(readWrite{}) { + t.Fatal("read-write DML must count as committed") + } + if !isCommittedMode(partitionedDML{}) { + t.Fatal("partitioned DML must count as committed") + } +} + +func TestMaterializeWithoutRows(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + o opts + want bool + }{ + {"default keeps rows", opts{}, false}, + {"redact drops rows", opts{RedactRows: true}, true}, + {"discard drops rows", opts{DiscardResults: true}, true}, + {"both drop rows", opts{RedactRows: true, DiscardResults: true}, true}, + } + for _, tc := range cases { + if got := materializeWithoutRows(tc.o); got != tc.want { + t.Errorf("%s: materializeWithoutRows = %v, want %v", tc.name, got, tc.want) + } + } +} + +func TestFinishPublishFailureAfterCommitIsWrapped(t *testing.T) { + t.Parallel() + + // Make the plan rename fail by turning the target path into a directory + // after the sinks opened their temp files. + dir := t.TempDir() + plan := filepath.Join(dir, "plan.json") + s, err := newOutputSinks(opts{ + Output: "-", + PlanOutput: plan, + QueryMode: "PROFILE", + }) + if err != nil { + t.Fatal(err) + } + if err := os.Mkdir(plan, 0o755); err != nil { + t.Fatal(err) + } + s.MarkPrimaryComplete() + finishErr := s.Finish(nil) + if finishErr == nil { + t.Fatal("expected publish failure") + } + // runCLI wraps this for committed modes; the wrapper must carry the + // no-rollback wording so a caller does not replay the DML. + got := wrapCommittedOutputError(finishErr).Error() + for _, want := range []string{"not a rollback", "not replayed", "plan output"} { + if !strings.Contains(got, want) { + t.Fatalf("wrapped error %q lacks %q", got, want) + } + } +} + +func TestProcessFlagsOutputDefaults(t *testing.T) { + oldArgs := os.Args + t.Cleanup(func() { os.Args = oldArgs }) + + os.Args = []string{"execspansql", "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"} + got, err := processFlags() + if err != nil { + t.Fatal(err) + } + if got.Output != "-" { + t.Fatalf("Output = %q, want -", got.Output) + } + if got.PlanOutput != "" || got.PlanFormat != "" || got.DiscardResults { + t.Fatalf("unexpected plan flags: %+v", got) + } + + os.Args = append([]string{"execspansql", "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"}, + "-o", "rows.json", "--plan-output", "plan.json", "--plan-format", "yaml", "--discard-results") + got, err = processFlags() + if err != nil { + t.Fatal(err) + } + if got.Output != "rows.json" || got.PlanOutput != "plan.json" || got.PlanFormat != "yaml" || !got.DiscardResults { + t.Fatalf("got %+v", got) + } +} + +func TestSplitModeValidationBeforeClient(t *testing.T) { + err := runMain(t, []string{ + "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1", + "--query-mode", "PROFILE", "--plan-output", "-", + }) + if err == nil || !strings.Contains(err.Error(), "stdout") { + t.Fatalf("both on stdout: error = %v", err) + } + + err = runMain(t, []string{ + "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1", + "--plan-output", "plan.json", + }) + if err == nil || !strings.Contains(err.Error(), "PLAN, PROFILE, or WITH_PLAN_AND_STATS") { + t.Fatalf("NORMAL plan-output: error = %v", err) + } +} + +type ioDiscardWriter struct{} + +func (ioDiscardWriter) Write([]byte) (int, error) { return 0, nil } + +func profileResultSetForSplitTest() *sppb.ResultSet { + return &sppb.ResultSet{ + Metadata: &sppb.ResultSetMetadata{ + RowType: &sppb.StructType{Fields: []*sppb.StructType_Field{ + {Name: "id", Type: &sppb.Type{Code: sppb.TypeCode_INT64}}, + }}, + }, + Rows: []*structpb.ListValue{{Values: []*structpb.Value{structpb.NewStringValue("1")}}}, + Stats: &sppb.ResultSetStats{ + QueryPlan: &sppb.QueryPlan{PlanNodes: []*sppb.PlanNode{{DisplayName: "Scan"}}}, + QueryStats: &structpb.Struct{Fields: map[string]*structpb.Value{ + "elapsed_time": structpb.NewStringValue("1 msecs"), + }}, + }, + } +} diff --git a/query_stats_modes_test.go b/query_stats_modes_test.go index 1a43fc7..44821dc 100644 --- a/query_stats_modes_test.go +++ b/query_stats_modes_test.go @@ -2,19 +2,23 @@ package main import ( "context" - "github.com/apstndb/execspansql/internal/grpctest" - "github.com/apstndb/spanemuboost" - "google.golang.org/api/option" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/status" + "encoding/csv" "net" + "os" + "path/filepath" "strings" "testing" "time" "cloud.google.com/go/spanner" sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/internal/grpctest" + "github.com/apstndb/spanemuboost" + "google.golang.org/api/option" "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/structpb" ) @@ -107,6 +111,7 @@ func TestQueryStatsResponseReachesOutput(t *testing.T) { type queryStatsModeServer struct { sppb.UnimplementedSpannerServer + omitValues bool } func (s *queryStatsModeServer) CreateSession(_ context.Context, req *sppb.CreateSessionRequest) (*sppb.Session, error) { @@ -114,13 +119,16 @@ func (s *queryStatsModeServer) CreateSession(_ context.Context, req *sppb.Create } func (s *queryStatsModeServer) ExecuteStreamingSql(_ *sppb.ExecuteSqlRequest, stream sppb.Spanner_ExecuteStreamingSqlServer) error { - if err := stream.Send(&sppb.PartialResultSet{ + first := &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 { + } + if !s.omitValues { + first.Values = []*structpb.Value{structpb.NewStringValue("value")} + } + if err := stream.Send(first); err != nil { return err } @@ -133,6 +141,160 @@ func (s *queryStatsModeServer) ExecuteStreamingSql(_ *sppb.ExecuteSqlRequest, st return stream.Send(&sppb.PartialResultSet{Stats: stats}) } +func startQueryStatsModeServer(t *testing.T, server *queryStatsModeServer) { + t.Helper() + 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()) +} + +func assertPlanEnvelopeFile(t *testing.T, planPath string) { + t.Helper() + b, err := os.ReadFile(planPath) + if err != nil { + t.Fatal(err) + } + var rs sppb.ResultSet + if err := protojson.Unmarshal(b, &rs); err != nil { + t.Fatalf("unmarshal plan: %v\n%s", err, b) + } + if rs.GetMetadata() == nil || rs.GetMetadata().GetRowType() == nil { + t.Fatalf("plan missing metadata: %s", b) + } + if len(rs.GetRows()) != 0 { + t.Fatalf("plan artifact has rows: %#v", rs.GetRows()) + } + if len(rs.GetStats().GetQueryPlan().GetPlanNodes()) == 0 { + t.Fatalf("plan missing queryPlan.planNodes: %s", b) + } +} + +func TestSplitPlanOutputCLI(t *testing.T) { + startQueryStatsModeServer(t, &queryStatsModeServer{}) + + t.Run("PROFILE csv", func(t *testing.T) { + dir := t.TempDir() + rowsPath := filepath.Join(dir, "rows.csv") + planPath := filepath.Join(dir, "plan.json") + err := runMain(t, []string{ + "database", "--project", "project", "--instance", "instance", + "--sql", "SELECT 'value'", + "--query-mode", "PROFILE", + "--format", "experimental_csv", + "--output", rowsPath, + "--plan-output", planPath, + "--timeout", "5s", + }) + if err != nil { + t.Fatal(err) + } + csvBytes, err := os.ReadFile(rowsPath) + if err != nil { + t.Fatal(err) + } + csvText := string(csvBytes) + if strings.Contains(csvText, "queryPlan") || strings.Contains(csvText, "Fake Scan") { + t.Fatalf("CSV contains plan: %s", csvText) + } + recs, err := csv.NewReader(strings.NewReader(csvText)).ReadAll() + if err != nil { + t.Fatal(err) + } + if len(recs) != 2 || recs[0][0] != "value" || recs[1][0] != "value" { + t.Fatalf("csv records = %#v", recs) + } + assertPlanEnvelopeFile(t, planPath) + }) + + t.Run("discard-results", func(t *testing.T) { + dir := t.TempDir() + rowsPath := filepath.Join(dir, "rows.csv") + planPath := filepath.Join(dir, "plan.json") + err := runMain(t, []string{ + "database", "--project", "project", "--instance", "instance", + "--sql", "SELECT 'value'", + "--query-mode", "PROFILE", + "--format", "experimental_csv", + "--output", rowsPath, + "--plan-output", planPath, + "--discard-results", + "--timeout", "5s", + }) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(rowsPath); !os.IsNotExist(err) { + t.Fatalf("primary file exists after --discard-results: %v", err) + } + assertPlanEnvelopeFile(t, planPath) + }) + + t.Run("lazy jq still writes plan after early stop", func(t *testing.T) { + dir := t.TempDir() + rowsPath := filepath.Join(dir, "rows.json") + planPath := filepath.Join(dir, "plan.json") + err := runMain(t, []string{ + "database", "--project", "project", "--instance", "instance", + "--sql", "SELECT 'value'", + "--query-mode", "PROFILE", + "--jq-input-mode", "lazy", + "--filter", ".rows[0]", + "--output", rowsPath, + "--plan-output", planPath, + "--timeout", "5s", + }) + if err != nil { + t.Fatal(err) + } + assertPlanEnvelopeFile(t, planPath) + }) +} + +func TestSplitPlanOutputZeroRowAndPLAN(t *testing.T) { + startQueryStatsModeServer(t, &queryStatsModeServer{omitValues: true}) + + for _, mode := range []string{"PROFILE", "PLAN"} { + t.Run(mode, func(t *testing.T) { + dir := t.TempDir() + rowsPath := filepath.Join(dir, "rows.csv") + planPath := filepath.Join(dir, "plan.json") + err := runMain(t, []string{ + "database", "--project", "project", "--instance", "instance", + "--sql", "SELECT 'value'", + "--query-mode", mode, + "--format", "experimental_csv", + "--output", rowsPath, + "--plan-output", planPath, + "--timeout", "5s", + }) + if err != nil { + t.Fatal(err) + } + csvBytes, err := os.ReadFile(rowsPath) + if err != nil { + t.Fatal(err) + } + recs, err := csv.NewReader(strings.NewReader(string(csvBytes))).ReadAll() + if err != nil { + t.Fatal(err) + } + if len(recs) != 1 || recs[0][0] != "value" { + t.Fatalf("csv records = %#v, want header only", recs) + } + assertPlanEnvelopeFile(t, planPath) + }) + } +} + // TestMainSendsAdditionalQueryStatsModes stops at the request boundary because // emulator support for these modes is independent of CLI option propagation. func TestMainSendsAdditionalQueryStatsModes(t *testing.T) {