Simplify command execution and result ownership

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.

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.

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

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

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

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