From ee765e667eb90a4677d274ac8b714668ea089250 Mon Sep 17 00:00:00 2001 From: apstndb <803393+apstndb@users.noreply.github.com> Date: Sat, 12 Sep 2026 15:53:06 +0900 Subject: [PATCH] Separate command preparation, execution, and result formatting --- AGENTS.md | 5 +- README.md | 4 + command.go | 88 ++++++++ command_test.go | 87 ++++++++ execution.go | 68 ++++++ execution_test.go | 117 +++++++++++ flags_validation_test.go | 7 +- format.go | 105 ++++++++++ integration_test.go | 33 +++ main.go | 427 ++++---------------------------------- output.go | 5 +- output_test.go | 25 +-- pdml_query_mode_test.go | 5 +- query_stats_modes_test.go | 2 +- reauth_test.go | 6 +- 15 files changed, 558 insertions(+), 426 deletions(-) create mode 100644 command.go create mode 100644 command_test.go create mode 100644 execution.go create mode 100644 execution_test.go create mode 100644 format.go diff --git a/AGENTS.md b/AGENTS.md index f68308c..444f884 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,7 +5,10 @@ `execspansql` is a Go CLI for Spanner query execution. It supports SQL parameter loading, multiple output formats (JSON/YAML/CSV), optional JQ filtering, and tracing options. Primary package layout: -- `main.go` / `trace.go` - CLI entrypoint, command wiring, execution flow. +- `main.go` / `trace.go` - CLI entrypoint, process lifecycle, client and tracing wiring. +- `command.go` - validated options, resolved SQL/parameters, and compiled jq. +- `execution.go` - transaction selection and owned query results; DML results become available only after commit. +- `format.go` / `output.go` - result formatting and destination publication; output errors never replay SQL. - `params/` - parameter file parsing and typed conversion helpers. - `resultset/` - Spanner result set materialization and formatting helpers. - `jqresult/` - JQ compile/execution pipeline and JSON conversion helpers. diff --git a/README.md b/README.md index 175e316..1ae12eb 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,8 @@ Document contents: 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. +DML result rows are buffered until commit before JSON, YAML, or CSV formatting begins. A transaction retry replaces the buffered result rather than emitting another copy. Redacted or discarded rows are not retained. Read-only CSV and lazy jq continue to consume rows incrementally. + 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. ``` @@ -259,6 +261,8 @@ In `lazy` mode, `metadata` is populated after the first row is read from Spanner 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. +Ctrl+C cancels both SQL execution and jq processing, including jq computations after SQL has completed. `--timeout` bounds SQL execution; it does not impose a separate deadline on subsequent jq processing. + #### Example: Extract QueryPlan [rendertree] command takes QueryPlan, and it can be extracted by jq filter. diff --git a/command.go b/command.go new file mode 100644 index 0000000..fbfb117 --- /dev/null +++ b/command.go @@ -0,0 +1,88 @@ +package main + +import ( + "fmt" + + "cloud.google.com/go/spanner" + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/jqresult" + "github.com/apstndb/execspansql/params" + "github.com/wader/gojq" +) + +// preparedCommand holds validated options and resolved inputs. Preparation does +// not authenticate, create a client, or open output files. +type preparedCommand struct { + opts + statement spanner.Statement + queryOptions spanner.QueryOptions + mode queryMode + jqMode jqresult.InputMode + jqCode *gojq.Code +} + +func prepareCommand(o opts) (*preparedCommand, error) { + jqMode, err := jqresult.ParseInputMode(o.JqInputMode) + if err != nil { + return nil, err + } + if err := jqMode.ValidateFormat(o.Format); err != nil { + return nil, err + } + if err := validateJqOutputOptions(o, jqMode); err != nil { + return nil, err + } + + var jqCode *gojq.Code + if !o.TryPartitionQuery && o.Format != "experimental_csv" { + jqFilter, err := readFileOrDefault(o.JqFromFile, o.JqFilter) + if err != nil { + return nil, err + } + if jqFilter == "" { + jqFilter = jqresult.DefaultFilter(jqMode) + } + + jqCode, err = jqresult.Compile(jqFilter, jqMode) + if err != nil { + return nil, err + } + } + + mode := sppb.ExecuteSqlRequest_QueryMode(sppb.ExecuteSqlRequest_QueryMode_value[o.QueryMode]) + queryOpts := queryOptionsFor(mode, o.Priority) + + query, err := readFileOrDefault(o.SqlFile, o.Sql) + if err != nil { + return nil, err + } + + tb, err := parseTimestampBound(o.TimestampBound.ReadTimestamp) + if err != nil { + return nil, fmt.Errorf("--read-timestamp is supplied but wrong: %w", err) + } + + m := queryModeForQuery(query, o.EnablePartitionedDML, tb) + if err := validateExecutionOptions(o, m); err != nil { + return nil, err + } + + // Freeze the statement (SQL and parameters) before any interactive step so + // a parameter file edited during a browser login cannot change what runs. + paramStrMap, err := o.mergedParams() + if err != nil { + return nil, err + } + paramMap, err := params.GenerateParams(paramStrMap, mode == sppb.ExecuteSqlRequest_PLAN) + if err != nil { + return nil, err + } + return &preparedCommand{ + opts: o, + statement: spanner.Statement{SQL: query, Params: paramMap}, + queryOptions: queryOpts, + mode: m, + jqMode: jqMode, + jqCode: jqCode, + }, nil +} diff --git a/command_test.go b/command_test.go new file mode 100644 index 0000000..24a22ba --- /dev/null +++ b/command_test.go @@ -0,0 +1,87 @@ +package main + +import ( + "bytes" + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/apstndb/execspansql/jqresult" +) + +func TestRunCLIInvalidArgumentsReturnsError(t *testing.T) { + if err := runCLI(t.Context(), []string{"--unknown-option"}); err == nil { + t.Fatal("expected argument error") + } +} + +func TestRunCLIHelpReturnsWithoutExecution(t *testing.T) { + out, err := captureStdout(t, func() error { return runCLI(t.Context(), []string{"--help"}) }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(out, "Usage: execspansql") { + t.Fatalf("help = %q", out) + } +} + +func TestPrepareCommandFreezesInputs(t *testing.T) { + dir := t.TempDir() + sql, params, filter := filepath.Join(dir, "query.sql"), filepath.Join(dir, "params.json"), filepath.Join(dir, "filter.jq") + for path, content := range map[string]string{sql: "SELECT @v", params: `{"v":1}`, filter: ".rows"} { + if err := os.WriteFile(path, []byte(content), 0600); err != nil { + t.Fatal(err) + } + } + o, err := processFlags([]string{"db", "--project", "p", "--instance", "i", "--sql-file", sql, "--param-file", params, "--filter-file", filter}) + if err != nil { + t.Fatal(err) + } + prepared, err := prepareCommand(o) + if err != nil { + t.Fatal(err) + } + for _, path := range []string{sql, params, filter} { + if err := os.Remove(path); err != nil { + t.Fatal(err) + } + } + if prepared.statement.SQL != "SELECT @v" || len(prepared.statement.Params) != 1 || prepared.jqCode == nil { + t.Fatalf("unresolved command: %+v", prepared) + } + value, ok := prepared.jqCode.Run(map[string]any{"rows": 7}).Next() + if !ok || value != 7 { + t.Fatalf("compiled filter result=%v, ok=%v", value, ok) + } +} + +func TestPrintJQHonorsCancellation(t *testing.T) { + for _, mode := range []jqresult.InputMode{jqresult.InputEager, jqresult.InputLazy} { + t.Run(string(mode), func(t *testing.T) { + code, err := jqresult.Compile("def spin: spin; spin", mode) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(t.Context()) + cancel() + var out bytes.Buffer + enc, err := newEncoder(&out, "json", false, false) + if err != nil { + t.Fatal(err) + } + var input any = map[string]any{} + if mode == jqresult.InputLazy { + input = jqresult.NewLazy(nil, false) + } + if err := printJQ(ctx, code, input, enc); !errors.Is(err, context.Canceled) { + t.Fatalf("error=%v, want canceled", err) + } + if out.Len() != 0 { + t.Fatalf("canceled filter emitted %q", out.String()) + } + }) + } +} diff --git a/execution.go b/execution.go new file mode 100644 index 0000000..da660be --- /dev/null +++ b/execution.go @@ -0,0 +1,68 @@ +package main + +import ( + "context" + "fmt" + + "cloud.google.com/go/spanner" + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/resultset" +) + +// queryResult owns either a live read-only iterator or a completed ResultSet. +// Write results are exposed only after commit, so output code never participates +// in transaction retries and cannot change the outcome of a successful write. +type queryResult struct { + rowIter *spanner.RowIterator + resultSet *sppb.ResultSet + committed bool +} + +func executeQuery(ctx context.Context, client *spanner.Client, command *preparedCommand) (*queryResult, error) { + result := &queryResult{} + switch mode := command.mode.(type) { + case single: + result.rowIter = client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, command.statement, command.queryOptions) + case readWrite: + _, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *spanner.ReadWriteTransaction) (err error) { + // Each attempt replaces the previous result. In particular, no rows from + // an aborted attempt can reach stdout or a file destination. + result.resultSet, err = resultset.Materialize(tx.QueryWithOptions(ctx, command.statement, command.queryOptions), + materializeWithoutRows(command.opts), spaniterStatsOpts(mode, command.queryOptions)...) + return err + }) + if err != nil { + return nil, err + } + result.committed = true + case partitionedDML: + count, err := client.PartitionedUpdateWithOptions(ctx, command.statement, command.queryOptions) + if err != nil { + return nil, err + } + result.resultSet = &sppb.ResultSet{ + Metadata: &sppb.ResultSetMetadata{RowType: &sppb.StructType{}}, + Stats: &sppb.ResultSetStats{RowCount: &sppb.ResultSetStats_RowCountLowerBound{RowCountLowerBound: count}}, + } + result.committed = true + default: + return nil, fmt.Errorf("unknown query mode: %T", mode) + } + return result, nil +} + +func (r *queryResult) materialize(redact bool) (*sppb.ResultSet, error) { + if r.resultSet != nil { + return r.resultSet, nil + } + var err error + r.resultSet, err = resultset.Materialize(r.rowIter, redact) + r.rowIter = nil // Materialize owns and stops the iterator, including on error. + return r.resultSet, err +} + +func (r *queryResult) Close() { + if r.rowIter != nil { + r.rowIter.Stop() + } +} diff --git a/execution_test.go b/execution_test.go new file mode 100644 index 0000000..c7e2ea9 --- /dev/null +++ b/execution_test.go @@ -0,0 +1,117 @@ +package main + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/structpb" + "google.golang.org/protobuf/types/known/timestamppb" +) + +// This transport fixture exercises the real SDK retry and output boundary. +// Database semantics remain covered by the emulator integration suite. +type executionServer struct { + queryStatsModeServer + executes atomic.Int32 + commits atomic.Int32 + retry bool + failCommit bool +} + +func (s *executionServer) BeginTransaction(context.Context, *sppb.BeginTransactionRequest) (*sppb.Transaction, error) { + return &sppb.Transaction{Id: []byte("test-transaction")}, nil +} + +func (s *executionServer) ExecuteStreamingSql(_ *sppb.ExecuteSqlRequest, stream sppb.Spanner_ExecuteStreamingSqlServer) error { + attempt := s.executes.Add(1) + return stream.Send(&sppb.PartialResultSet{ + Metadata: &sppb.ResultSetMetadata{ + Transaction: &sppb.Transaction{Id: []byte("test-transaction")}, + RowType: &sppb.StructType{Fields: []*sppb.StructType_Field{{Name: "value", Type: &sppb.Type{Code: sppb.TypeCode_STRING}}}}, + }, + Values: []*structpb.Value{structpb.NewStringValue(fmt.Sprintf("attempt-%d", attempt))}, + Stats: &sppb.ResultSetStats{RowCount: &sppb.ResultSetStats_RowCountExact{RowCountExact: 1}}, + }) +} + +func (s *executionServer) Commit(context.Context, *sppb.CommitRequest) (*sppb.CommitResponse, error) { + attempt := s.commits.Add(1) + if s.failCommit { + return nil, status.Error(codes.FailedPrecondition, "test commit failure") + } + if s.retry && attempt == 1 { + return nil, status.Error(codes.Aborted, "retry transaction") + } + return &sppb.CommitResponse{CommitTimestamp: timestamppb.Now()}, nil +} + +func (*executionServer) Rollback(context.Context, *sppb.RollbackRequest) (*emptypb.Empty, error) { + return &emptypb.Empty{}, nil +} + +func TestDMLResultPublication(t *testing.T) { + for _, format := range []string{"json", "yaml", "experimental_csv"} { + for _, scenario := range []string{"retry", "commit_failure", "output_failure"} { + t.Run(format+"/"+scenario, func(t *testing.T) { + server := &executionServer{retry: scenario == "retry", failCommit: scenario == "commit_failure"} + startQueryStatsModeServer(t, server) + path := filepath.Join(t.TempDir(), "result") + if scenario == "output_failure" { + if err := os.Mkdir(path, 0700); err != nil { + t.Fatal(err) + } + } else if err := os.WriteFile(path, []byte("original"), 0600); err != nil { + t.Fatal(err) + } + err := runCLI(t.Context(), []string{"db", "--project", "p", "--instance", "i", + "--sql", "UPDATE T SET V=1 THEN RETURN V", "--format", format, "--output", path, "--timeout", "5s"}) + switch scenario { + case "retry": + if err != nil { + t.Fatal(err) + } + output, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(output), "attempt-1") || strings.Count(string(output), "attempt-2") != 1 { + t.Fatalf("output must contain only the committed attempt: %s", output) + } + if got := server.executes.Load(); got != 2 { + t.Fatalf("executions=%d, want 2", got) + } + case "commit_failure": + if err == nil || !strings.Contains(err.Error(), "test commit failure") || strings.Contains(err.Error(), "statement was committed") { + t.Fatalf("commit error = %v", err) + } + output, readErr := os.ReadFile(path) + if readErr != nil { + t.Fatal(readErr) + } + if string(output) != "original" { + t.Fatalf("failed commit published %q", output) + } + case "output_failure": + if err == nil || !strings.Contains(err.Error(), "statement was committed") { + t.Fatalf("output error = %v", err) + } + if got := server.executes.Load(); got != 1 { + t.Fatalf("output failure replayed SQL: executions=%d", got) + } + if got := server.commits.Load(); got != 1 { + t.Fatalf("commits=%d, want 1", got) + } + } + }) + } + } +} diff --git a/flags_validation_test.go b/flags_validation_test.go index 92ddbdb..496bb54 100644 --- a/flags_validation_test.go +++ b/flags_validation_test.go @@ -1,7 +1,6 @@ package main import ( - "os" "strings" "testing" "time" @@ -537,9 +536,6 @@ func TestQueryOptionsForPriority(t *testing.T) { } func TestProcessFlagsPriority(t *testing.T) { - oldArgs := os.Args - t.Cleanup(func() { os.Args = oldArgs }) - baseArgs := []string{"execspansql", "database", "--project", "project", "--instance", "instance", "--sql", "SELECT 1"} tests := []struct { name string @@ -557,8 +553,7 @@ func TestProcessFlagsPriority(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - os.Args = tt.args - got, err := processFlags() + got, err := processFlags(tt.args[1:]) if tt.wantErr != "" { if err == nil || !strings.Contains(err.Error(), tt.wantErr) { t.Fatalf("processFlags() error = %v, want %q", err, tt.wantErr) diff --git a/format.go b/format.go new file mode 100644 index 0000000..c6743e2 --- /dev/null +++ b/format.go @@ -0,0 +1,105 @@ +package main + +import ( + "context" + + sppb "cloud.google.com/go/spanner/apiv1/spannerpb" + "github.com/apstndb/execspansql/jqresult" + "github.com/wader/gojq" +) + +// writeResult chooses only an output strategy. Transaction selection, retries, +// and commit handling belong to executeQuery and runCLI. +func (c *preparedCommand) writeResult(ctx context.Context, result *queryResult, sinks *outputSinks) error { + // Discarded results need metadata/stats only. DML is already materialized; + // read-only CSV and lazy jq retain their streaming behavior. + if c.DiscardResults || result.resultSet != nil || (c.Format != "experimental_csv" && c.jqMode == jqresult.InputEager) { + rs, err := result.materialize(materializeWithoutRows(c.opts)) + if err != nil { + return err + } + return c.writeResultSet(ctx, rs, sinks) + } + + if c.Format == "experimental_csv" { + result, err := writeCsvFromRowIter(sinks.primary, result.rowIter, c.RedactRows) + if err != nil { + return err + } + sinks.MarkPrimaryComplete() + if !sinks.hasPlan { + return nil + } + stats, err := statsFromWriterResult(result) + if err != nil { + return err + } + return writePlan(sinks.plan, effectivePlanFormat(c.opts), result.Metadata, stats) + } + + enc, err := newEncoder(sinks.primary, c.Format, c.CompactOutput, c.JqRawOutput) + if err != nil { + return err + } + var lazyOpts []jqresult.LazyOption + if sinks.hasPlan { + lazyOpts = append(lazyOpts, jqresult.WithOmitQueryPlan()) + } + lazy := jqresult.NewLazy(result.rowIter, c.RedactRows, lazyOpts...) + defer lazy.Stop() + if err := printJQ(ctx, c.jqCode, lazy, enc); err != nil { + return err + } + sinks.MarkPrimaryComplete() + if !sinks.hasPlan { + return nil + } + if err := lazy.Drain(); err != nil { + return err + } + drained := lazy.Result() + stats, err := drained.StatsProto() + if err != nil { + return err + } + return writePlan(sinks.plan, effectivePlanFormat(c.opts), drained.Metadata, stats) +} + +func (c *preparedCommand) writeResultSet(ctx context.Context, rs *sppb.ResultSet, sinks *outputSinks) error { + stats, metadata := rs.Stats, rs.Metadata + if sinks.hasPlan { + stats, metadata = stripQueryPlanForPrimary(rs) + } + if sinks.primary != nil { + if c.Format == "experimental_csv" { + if err := writeCsvFromResultSet(sinks.primary, rs); err != nil { + return err + } + } else { + input, err := jqresult.ResultSetMap(rs) + if err != nil { + return err + } + enc, err := newEncoder(sinks.primary, c.Format, c.CompactOutput, c.JqRawOutput) + if err != nil { + return err + } + if err := printJQ(ctx, c.jqCode, input, enc); err != nil { + return err + } + } + } + sinks.MarkPrimaryComplete() + return writePlan(sinks.plan, effectivePlanFormat(c.opts), metadata, stats) +} + +// printJQ owns encoder completion on both success and failure. Use the caller's +// cancellation context even after SQL has completed and no RPC remains active. +func printJQ(ctx context.Context, code *gojq.Code, input any, enc encoder) (err error) { + defer func() { + if closeErr := closeEncoder(enc); err == nil { + err = closeErr + } + }() + return jqresult.Print(enc, code.RunWithContext(ctx, input)) +} diff --git a/integration_test.go b/integration_test.go index 8c6b756..8ab3039 100644 --- a/integration_test.go +++ b/integration_test.go @@ -648,6 +648,39 @@ func TestWithCloudSpannerEmulator(t *testing.T) { }) }) + t.Run("CLI DML formats publish committed rows", func(t *testing.T) { + t.Setenv("SPANNER_EMULATOR_HOST", env.Emulator().URI()) + for _, format := range []string{"json", "yaml", "experimental_csv"} { + t.Run(format, func(t *testing.T) { + const value = "committed-cli-result" + out, err := captureStdout(t, func() error { + return runMain(t, []string{env.DatabaseID, "--project", env.ProjectID, "--instance", env.InstanceID, + "--sql", "UPDATE Singers SET FirstName='" + value + "' WHERE SingerId=1 THEN RETURN FirstName", "--format", format}) + }) + if err != nil { + t.Fatal(err) + } + if strings.Count(out, value) != 1 { + t.Fatalf("expected one committed row: %s", out) + } + if format != "experimental_csv" && !strings.Contains(out, "rowCountExact") { + t.Fatalf("DML count missing: %s", out) + } + row, err := client.Single().ReadRow(ctx, "Singers", spanner.Key{1}, []string{"FirstName"}) + if err != nil { + t.Fatal(err) + } + var stored string + if err := row.Column(0, &stored); err != nil { + t.Fatal(err) + } + if stored != value { + t.Fatalf("stored=%q, want %q", stored, value) + } + }) + } + }) + t.Run("split plan output", func(t *testing.T) { t.Setenv("SPANNER_EMULATOR_HOST", env.Emulator().URI()) diff --git a/main.go b/main.go index b632625..7ab0dd3 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,6 @@ package main import ( - "bytes" "context" "errors" "io" @@ -28,12 +27,10 @@ import ( "github.com/alecthomas/kong" "github.com/apstndb/execspansql/jqresult" "github.com/apstndb/execspansql/params" - "github.com/apstndb/execspansql/resultset" "github.com/apstndb/gsqlutils/stmtkind" "github.com/apstndb/spaniter" "github.com/apstndb/spannerotel/interceptor" svwriter "github.com/apstndb/spanvalue/writer" - "github.com/wader/gojq" ) const ( @@ -43,7 +40,14 @@ const ( ) func main() { - if err := _main(); err != nil { + // Keep process-wide signals and exit handling outside the testable runner. + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + go func() { + <-ctx.Done() + stop() // A second interrupt uses the default handler. + }() + if err := runCLI(ctx, os.Args[1:]); err != nil { log.Fatalln(err) } } @@ -131,10 +135,18 @@ func (o opts) mergedParams() (map[string]string, error) { return params.MergeParams(fileParams, cliParams), nil } -func processFlags() (o opts, err error) { +var errHelpRequested = errors.New("help requested") + +func processFlags(args []string) (o opts, err error) { parser, err := kong.New(&o, kong.Name("execspansql"), kong.Description("Yet another gcloud spanner databases execute-sql replacement"), + kong.Help(func(options kong.HelpOptions, ctx *kong.Context) error { + if err := kong.DefaultHelpPrinter(options, ctx); err != nil { + return err + } + return errHelpRequested // End parsing without Kong calling os.Exit. + }), kong.ExplicitGroups([]kong.Group{ {Key: "Timestamp Bound", Title: "Timestamp Bound"}, }), @@ -142,12 +154,10 @@ func processFlags() (o opts, err error) { if err != nil { return o, err } - defer func() { - if err != nil { - fmt.Fprintln(os.Stderr, "error:", err) - } - }() - ctx, err := parser.Parse(os.Args[1:]) + ctx, err := parser.Parse(args) + if errors.Is(err, errHelpRequested) { + return o, err + } if err != nil { var parseErr *kong.ParseError if errors.As(err, &parseErr) { @@ -383,115 +393,22 @@ func spaniterStatsOpts(mode queryMode, opts spanner.QueryOptions) []spaniter.Opt return nil } -func runInNewTransaction(ctx context.Context, client *spanner.Client, stmt spanner.Statement, opts spanner.QueryOptions, mode queryMode, reductRows bool) (*sppb.ResultSet, error) { - statOpts := spaniterStatsOpts(mode, opts) - var rs *sppb.ResultSet - switch mode := mode.(type) { - case readWrite: - _, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *spanner.ReadWriteTransaction) (err error) { - rs, err = resultset.Materialize(tx.QueryWithOptions(ctx, stmt, opts), reductRows, statOpts...) - return err - }) - return rs, err - case single: - return resultset.Materialize(client.Single().WithTimestampBound(mode.TimestampBound).QueryWithOptions(ctx, stmt, opts), reductRows, statOpts...) - case partitionedDML: - count, err := client.PartitionedUpdateWithOptions(ctx, stmt, opts) - return &sppb.ResultSet{ - Metadata: &sppb.ResultSetMetadata{ - RowType: &sppb.StructType{}, - }, - Stats: &sppb.ResultSetStats{ - RowCount: &sppb.ResultSetStats_RowCountLowerBound{RowCountLowerBound: count}, - }, - }, err - default: - panic(fmt.Sprintf("unknown mode: %T", mode)) - } -} - -func _main() error { - return runCLI() -} - // runCLI accepts client options so transport tests can inspect outgoing RPCs. // Non-empty clientOptions skip the ADC reauth preflight (tests inject insecure // dial options that bypass application-default credentials). -func runCLI(clientOptions ...option.ClientOption) (err error) { - o, err := processFlags() - if err != nil { - os.Exit(1) - } - defer func() { err = wrapWithHint(err) }() - - // The first interrupt cancels ctx so an in-progress gcloud login or query - // unwinds cleanly; stop() then restores default signal handling so a - // second interrupt still terminates the process if shutdown hangs. - ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) - defer stop() - go func() { - <-ctx.Done() - stop() - }() - - jqMode, err := jqresult.ParseInputMode(o.JqInputMode) - if err != nil { - return err - } - if err := jqMode.ValidateFormat(o.Format); err != nil { - return err - } - if err := validateJqOutputOptions(o, jqMode); err != nil { - return err - } - - var ( - jqCode *gojq.Code - ) - if !o.TryPartitionQuery && o.Format != "experimental_csv" { - jqFilter, err := readFileOrDefault(o.JqFromFile, o.JqFilter) - if err != nil { - return err - } - if jqFilter == "" { - jqFilter = jqresult.DefaultFilter(jqMode) - } - - jqCode, err = jqresult.Compile(jqFilter, jqMode) - if err != nil { - return err - } - } - - mode := sppb.ExecuteSqlRequest_QueryMode(sppb.ExecuteSqlRequest_QueryMode_value[o.QueryMode]) - queryOpts := queryOptionsFor(mode, o.Priority) - - query, err := readFileOrDefault(o.SqlFile, o.Sql) - if err != nil { - return err +func runCLI(ctx context.Context, args []string, clientOptions ...option.ClientOption) (err error) { + o, err := processFlags(args) + if errors.Is(err, errHelpRequested) { + return nil } - - tb, err := parseTimestampBound(o.TimestampBound.ReadTimestamp) if err != nil { - return fmt.Errorf("--read-timestamp is supplied but wrong: %w", err) - } - - m := queryModeForQuery(query, o.EnablePartitionedDML, tb) - if err := validateExecutionOptions(o, m); err != nil { return err } - - // Freeze the statement (SQL and parameters) before any interactive step so - // a parameter file edited during a browser login cannot change what runs. - paramStrMap, err := o.mergedParams() - if err != nil { - return err - } - paramMap, err := params.GenerateParams(paramStrMap, mode == sppb.ExecuteSqlRequest_PLAN) + defer func() { err = wrapWithHint(err) }() + command, err := prepareCommand(o) if err != nil { return err } - stmt := spanner.Statement{SQL: query, Params: paramMap} sinks, err := newOutputSinks(o) if err != nil { @@ -506,6 +423,7 @@ func runCLI(clientOptions ...option.ClientOption) (err error) { // Query execution timeout starts after authentication. Login is // human-paced and must not consume --timeout. + outputCtx := ctx // --timeout bounds SQL; interrupts also cancel later jq processing. ctx, cancel := context.WithTimeout(ctx, o.Timeout) defer cancel() @@ -528,14 +446,14 @@ func runCLI(clientOptions ...option.ClientOption) (err error) { defer client.Close() if o.TryPartitionQuery { - bt, err := client.BatchReadOnlyTransaction(ctx, tb) + bt, err := client.BatchReadOnlyTransaction(ctx, command.mode.(single).TimestampBound) if err != nil { return err } defer bt.Close() defer func() { bt.Cleanup(ctx) }() - _, err = bt.PartitionQuery(ctx, stmt, spanner.PartitionOptions{}) + _, err = bt.PartitionQuery(ctx, command.statement, spanner.PartitionOptions{}) if err != nil { return err } @@ -549,31 +467,19 @@ func runCLI(clientOptions ...option.ClientOption) (err error) { return sinks.Finish(nil) } - var workErr error - if o.Format == "experimental_csv" { - 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 -} + result, err := executeQuery(ctx, client, command) + if err != nil { + return err + } + defer result.Close() -// 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 + // Every format shares the same commit and publication boundary. No output + // error can re-enter the transaction or cause the statement to be replayed. + err = sinks.Finish(command.writeResult(outputCtx, result, sinks)) + if err != nil && result.committed { + return wrapCommittedOutputError(err) } + return err } // materializeWithoutRows reports whether the eager path may drop row values @@ -583,117 +489,6 @@ 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, 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() - 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 - } - 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: - 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 - } - 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)) - } -} - // csvRedactRowIteratorWriter implements [svwriter.RowIteratorWriter] for --redact-rows CSV: // it registers schema and flushes the header via the embedded [svwriter.DelimitedWriter] but // discards row bodies in WriteRow while WriteRowIterator drains the iterator. @@ -704,7 +499,7 @@ type csvRedactRowIteratorWriter struct { 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. +// WriteRowIterator stops the iterator; queryResult also closes it on early failures. func writeCsvFromRowIter(writer io.Writer, rowIter *spanner.RowIterator, redactRows bool) (*svwriter.RowIteratorResult, error) { csvWriter, err := svwriter.NewCSVWriter(writer) if err != nil { @@ -724,8 +519,7 @@ func prepareCsvRowType(csvWriter *svwriter.DelimitedWriter, metadata *sppb.Resul return csvWriter.PrepareRowType(metadata.GetRowType()) } -// writeCsvFromResultSet writes CSV from an in-memory ResultSet. Used by unit tests -// and partitioned DML (no RowIterator). WithMetadata at construction is appropriate here. +// writeCsvFromResultSet writes completed DML results without a live iterator. func writeCsvFromResultSet(writer io.Writer, rs *sppb.ResultSet) error { if rs == nil || rs.GetMetadata() == nil || rs.GetMetadata().GetRowType() == nil { return errors.New("result set metadata is missing or invalid") @@ -791,141 +585,6 @@ func closeEncoder(enc encoder) error { return nil } -func runJqOutput( - ctx context.Context, - client *spanner.Client, - stmt spanner.Statement, - opts spanner.QueryOptions, - mode queryMode, - 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, materializeWithoutRows(o)) - 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 - } - 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) - } - } - sinks.MarkPrimaryComplete() - if sinks.plan != nil { - return wrap(writePlan(sinks.plan, planFmt, metadata, planStats)) - } - return nil - } - - switch mode := mode.(type) { - case single: - 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 - } - 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: - panic(fmt.Sprintf("unknown mode: %T", mode)) - } -} - -func runJqOnRowIter( - rowIter *spanner.RowIterator, - redactRows bool, - jqCode *gojq.Code, - enc encoder, - sinks *outputSinks, - planFmt string, -) error { - 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 - } - return writePlan(sinks.plan, planFmt, result.Metadata, stats) -} - func newEncoder(writer io.Writer, format string, compactOutput bool, rawOutput bool) (encoder, error) { switch format { case "yaml": diff --git a/output.go b/output.go index 5a56a29..8419351 100644 --- a/output.go +++ b/output.go @@ -458,7 +458,7 @@ func writePlan(w io.Writer, format string, metadata *sppb.ResultSetMetadata, sta return closeEncoder(enc) } -func statsFromWriterResult(r *svwriter.RowIteratorResult, encodeRowCount bool) (*sppb.ResultSetStats, error) { +func statsFromWriterResult(r *svwriter.RowIteratorResult) (*sppb.ResultSetStats, error) { if r == nil { return nil, nil } @@ -472,9 +472,6 @@ func statsFromWriterResult(r *svwriter.RowIteratorResult, encodeRowCount bool) ( } 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 } diff --git a/output_test.go b/output_test.go index 18dc567..aacdea8 100644 --- a/output_test.go +++ b/output_test.go @@ -437,20 +437,6 @@ func TestDiscardResultsProducesNoPrimaryBytes(t *testing.T) { } } -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() @@ -505,11 +491,8 @@ func TestFinishPublishFailureAfterCommitIsWrapped(t *testing.T) { } 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() + args := []string{"database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"} + got, err := processFlags(args) if err != nil { t.Fatal(err) } @@ -520,9 +503,9 @@ func TestProcessFlagsOutputDefaults(t *testing.T) { t.Fatalf("unexpected plan flags: %+v", got) } - os.Args = append([]string{"execspansql", "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"}, + args = append([]string{"database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"}, "-o", "rows.json", "--plan-output", "plan.json", "--plan-format", "yaml", "--discard-results") - got, err = processFlags() + got, err = processFlags(args) if err != nil { t.Fatal(err) } diff --git a/pdml_query_mode_test.go b/pdml_query_mode_test.go index e845579..0c12ac7 100644 --- a/pdml_query_mode_test.go +++ b/pdml_query_mode_test.go @@ -76,10 +76,7 @@ func TestPartitionedDMLQueryMode(t *testing.T) { func runMain(t *testing.T, args []string, options ...option.ClientOption) error { t.Helper() - old := os.Args - os.Args = append([]string{"execspansql"}, args...) - defer func() { os.Args = old }() - return runCLI(options...) + return runCLI(t.Context(), args, options...) } func readPdmlQueryModeV(t *testing.T, ctx context.Context, client *spanner.Client) int64 { diff --git a/query_stats_modes_test.go b/query_stats_modes_test.go index 44821dc..c730d1d 100644 --- a/query_stats_modes_test.go +++ b/query_stats_modes_test.go @@ -141,7 +141,7 @@ func (s *queryStatsModeServer) ExecuteStreamingSql(_ *sppb.ExecuteSqlRequest, st return stream.Send(&sppb.PartialResultSet{Stats: stats}) } -func startQueryStatsModeServer(t *testing.T, server *queryStatsModeServer) { +func startQueryStatsModeServer(t *testing.T, server sppb.SpannerServer) { t.Helper() listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { diff --git a/reauth_test.go b/reauth_test.go index 296429e..51cb68e 100644 --- a/reauth_test.go +++ b/reauth_test.go @@ -154,9 +154,6 @@ func TestWrapWithHintMentionsEnv(t *testing.T) { } func TestProcessFlagsReauth(t *testing.T) { - oldArgs := os.Args - t.Cleanup(func() { os.Args = oldArgs }) - base := []string{"execspansql", "database", "--project", "p", "--instance", "i", "--sql", "SELECT 1"} tests := []struct { name string @@ -179,8 +176,7 @@ func TestProcessFlagsReauth(t *testing.T) { t.Setenv("EXECSPANSQL_REAUTH", "") _ = os.Unsetenv("EXECSPANSQL_REAUTH") } - os.Args = tt.args - got, err := processFlags() + got, err := processFlags(tt.args[1:]) if tt.wantErr != "" { if err == nil || !strings.Contains(err.Error(), tt.wantErr) { t.Fatalf("processFlags() error = %v, want %q", err, tt.wantErr)