Fix CLI bugs: DML classification, trace timeout, partition timestamp

.github/workflows/ko.yml

@@ -10,12 +10,12 @@ jobs:
     name: Publish
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/setup-go@v4
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
         with:
-          go-version: '1.24.x'
-      - uses: actions/checkout@v3
+          go-version-file: go.mod
 
       - uses: ko-build/setup-ko@v0.6
       - run: |
           tag=$(echo ${{ github.ref }} | cut -c11-)
-          ko build -t latest,${tag} -B
+          ko build -t latest,${tag} -B

README.md

@@ -55,7 +55,7 @@ Arguments:
   database:                                    (required) ID of the database.
 ```
 
-Local build requires Go 1.24.
+Local build requires Go 1.25.
 
 ```
 $ go install github.com/apstndb/execspansql@latest

context_test.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestWithCancelPreservesParentDeadline(t *testing.T) {
+	t.Parallel()
+
+	parent, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	child, childCancel := context.WithCancel(parent)
+	defer childCancel()
+
+	got, ok := child.Deadline()
+	if !ok {
+		t.Fatal("child deadline missing")
+	}
+	want, _ := parent.Deadline()
+	if !got.Equal(want) {
+		t.Fatalf("deadline = %v, want %v", got, want)
+	}
+}

go.mod

@@ -6,12 +6,13 @@ require (
 	cloud.google.com/go/spanner v1.84.1
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.24.0
 	github.com/alecthomas/kong v1.15.0
-	github.com/apstndb/gsqlsep v0.0.0-20240823174243-432be37d515a
+	github.com/apstndb/gsqlutils v0.0.0-20260502161854-d7d6011a36e0
 	github.com/apstndb/memebridge v0.6.1
 	github.com/apstndb/spanemuboost v0.4.6
 	github.com/apstndb/spaniter v0.3.0
 	github.com/apstndb/spannerotel v0.0.0-20220113012943-45b1c9cc5afb
 	github.com/apstndb/spanvalue v0.8.0
+	github.com/cloudspannerecosystem/memefish v0.6.2
 	github.com/goccy/go-yaml v1.19.2
 	github.com/google/go-cmp v0.7.0
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
@@ -44,7 +45,6 @@ require (
 	github.com/apstndb/spantype v0.3.11 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
-	github.com/cloudspannerecosystem/memefish v0.6.2 // indirect
 	github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
@@ -107,7 +107,6 @@ require (
 	go.opentelemetry.io/otel/trace v1.41.0 // indirect
 	go.uber.org/multierr v1.10.0 // indirect
 	golang.org/x/crypto v0.48.0 // indirect
-	golang.org/x/exp v0.0.0-20240823005443-9b4947da3948 // indirect
 	golang.org/x/net v0.49.0 // indirect
 	golang.org/x/oauth2 v0.34.0 // indirect
 	golang.org/x/sync v0.21.0 // indirect

go.sum

@@ -656,8 +656,8 @@ github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kd
 github.com/apache/arrow/go/v10 v10.0.1/go.mod h1:YvhnlEePVnBS4+0z3fhPfUy7W1Ikj0Ih0vcRo/gZ1M0=
 github.com/apache/arrow/go/v11 v11.0.0/go.mod h1:Eg5OsL5H+e299f7u5ssuXsuHQVEGC4xei5aX110hRiI=
 github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU=
-github.com/apstndb/gsqlsep v0.0.0-20240823174243-432be37d515a h1:9rutREITQVQSkSt2uoRr8ap6JA+k5i7NL1zRmJe621g=
-github.com/apstndb/gsqlsep v0.0.0-20240823174243-432be37d515a/go.mod h1:NQogaK8AOkyzXEXMqGvc6hV0SKO8cUkcqqXJimyvTlw=
+github.com/apstndb/gsqlutils v0.0.0-20260502161854-d7d6011a36e0 h1:VySrhGRfqXUDCrUuAEz70DrDnPrvj/nXGnUHwJF3PG4=
+github.com/apstndb/gsqlutils v0.0.0-20260502161854-d7d6011a36e0/go.mod h1:VwhJDip+HamDWWt+Ol4ECFU0g0HiveA27DeNkt3U8S0=
 github.com/apstndb/memebridge v0.6.1 h1:M1FMF5kb5vP/v5rTfP/N1HlLCfBPX7ctDQ3K/gFkI6c=
 github.com/apstndb/memebridge v0.6.1/go.mod h1:E/HVP4iaSgiPXMIQTPT1u1uKkjmRtRECkVKE2TrF3bc=
 github.com/apstndb/spanemuboost v0.4.6 h1:9f1qQDLNrPQJiswNLtiTaR0U//zToqxjcEGWGkyThm4=
@@ -917,6 +917,7 @@ github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
 github.com/jung-kurt/gofpdf v1.0.0/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=
 github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=
+github.com/k0kubun/pp v1.3.1-0.20200204103551-99835366d1cc h1:XLjmW07gT7cG/wb6mavIrvAIWBYaTacPo8UOnxGSspA=
 github.com/k0kubun/pp/v3 v3.4.1 h1:1WdFZDRRqe8UsR61N/2RoOZ3ziTEqgTPVqKrHeb779Y=
 github.com/k0kubun/pp/v3 v3.4.1/go.mod h1:+SiNiqKnBfw1Nkj82Lh5bIeKQOAkPy6Xw9CAZUZ8npI=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
@@ -1125,8 +1126,6 @@ golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u0
 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM=
 golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU=
 golang.org/x/exp v0.0.0-20220827204233-334a2380cb91/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
-golang.org/x/exp v0.0.0-20240823005443-9b4947da3948 h1:kx6Ds3MlpiUHKj7syVnbp57++8WpuKPcR5yjLBjvLEA=
-golang.org/x/exp v0.0.0-20240823005443-9b4947da3948/go.mod h1:akd2r19cwCdwSwWeIdzYQGa/EZZyqcOdwWiwj5L5eKQ=
 golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs=
 golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
 golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0=

integration_test.go

@@ -12,9 +12,9 @@ import (
 
 	"cloud.google.com/go/spanner"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
-	"github.com/apstndb/gsqlsep"
 	"github.com/apstndb/spanemuboost"
 	"github.com/apstndb/spaniter"
+	"github.com/cloudspannerecosystem/memefish"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/structpb"
@@ -30,6 +30,29 @@ var ddl string
 //go:embed testdata/dml.sql
 var dml string
 
+func splitSQLStatements(s string) ([]string, error) {
+	stmts, err := memefish.SplitRawStatements("", s)
+	if err != nil {
+		return nil, err
+	}
+	out := make([]string, 0, len(stmts))
+	for _, stmt := range stmts {
+		if trimmed := strings.TrimSpace(stmt.Statement); trimmed != "" {
+			out = append(out, trimmed)
+		}
+	}
+	return out, nil
+}
+
+func mustSplitSQLStatements(t *testing.T, s string) []string {
+	t.Helper()
+	out, err := splitSQLStatements(s)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return out
+}
+
 func collectRows(seq iter.Seq2[*spanner.Row, error]) ([]*spanner.Row, error) {
 	var rows []*spanner.Row
 	for row, err := range seq {
@@ -82,8 +105,8 @@ func TestWithCloudSpannerEmulator(t *testing.T) {
 	ctx := context.Background()
 
 	env, err := spanemuboost.RunEmulatorWithClients(ctx,
-		spanemuboost.WithSetupDDLs(gsqlsep.SeparateInputString(ddl)),
-		spanemuboost.WithSetupRawDMLs(gsqlsep.SeparateInputString(dml)),
+		spanemuboost.WithSetupDDLs(mustSplitSQLStatements(t, ddl)),
+		spanemuboost.WithSetupRawDMLs(mustSplitSQLStatements(t, dml)),
 	)
 	if err != nil {
 		t.Fatal(err)

main.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"github.com/apstndb/execspansql/params"
 	"io"
-	"regexp"
 	"time"
 
 	"fmt"
@@ -34,6 +33,7 @@ import (
 	"github.com/alecthomas/kong"
 	"github.com/apstndb/execspansql/jqresult"
 	"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"
@@ -46,8 +46,6 @@ func main() {
 	}
 }
 
-var dmlRe = regexp.MustCompile(`(?is)^\s*(INSERT|UPDATE|DELETE)\s+.+$`)
-
 var debuglog *log.Logger
 
 func init() {
@@ -302,18 +300,17 @@ func _main() error {
 		otel.SetTracerProvider(tp)
 		oteloc.InstallTraceBridge()
 
-		ctx, cancel := context.WithCancel(context.Background())
-		defer cancel()
+		traceCtx, traceCancel := context.WithCancel(ctx)
+		defer traceCancel()
+		ctx = traceCtx
 
-		// Cleanly shutdown and flush telemetry when the application exits.
-		defer func(ctx context.Context) {
-			// Do not make the application hang when it is shutdown.
-			ctx, cancel = context.WithTimeout(ctx, time.Second*5)
-			defer cancel()
-			if err := tp.Shutdown(ctx); err != nil {
-				log.Fatal(err)
+		defer func() {
+			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer shutdownCancel()
+			if err := tp.Shutdown(shutdownCtx); err != nil {
+				log.Printf("trace provider shutdown: %v", err)
 			}
-		}(ctx)
+		}()
 	}
 
 	logGrpc := o.LogGrpc
@@ -348,7 +345,7 @@ func _main() error {
 	switch {
 	case o.EnablePartitionedDML:
 		m = partitionedDML{}
-	case dmlRe.MatchString(query):
+	case stmtkind.IsDMLLexical(query):
 		m = readWrite{}
 	default:
 		m = single{tb}
@@ -357,18 +354,18 @@ func _main() error {
 	stmt := spanner.Statement{SQL: query, Params: paramMap}
 
 	if o.TryPartitionQuery {
-		bt, err := client.BatchReadOnlyTransaction(ctx, spanner.StrongRead())
+		bt, err := client.BatchReadOnlyTransaction(ctx, tb)
 		if err != nil {
 			return err
 		}
 		defer bt.Close()
+		defer func() { bt.Cleanup(ctx) }()
 
 		_, err = bt.PartitionQuery(ctx, stmt, spanner.PartitionOptions{})
 		if err != nil {
 			return err
 		}
 
-		bt.Cleanup(ctx)
 		fmt.Println("success")
 		return nil
 	}