From adaa384b599544cf52422b5fde41be48c362a0ce Mon Sep 17 00:00:00 2001
From: apstndb <803393+apstndb@users.noreply.github.com>
Date: Fri, 10 Jul 2026 04:10:56 +0900
Subject: [PATCH] fix: validate partitioned DML before client creation

---
 flags_validation_test.go | 36 ++++++++++++++++++++++++++++++++++++
 main.go                  | 14 +++++++++++---
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/flags_validation_test.go b/flags_validation_test.go
index 6aef1a8..918d8ed 100644
--- a/flags_validation_test.go
+++ b/flags_validation_test.go
@@ -266,6 +266,28 @@ func TestValidateExecutionOptions(t *testing.T) {
 			mode: partitionedDML{},
 			err:  "--strong cannot be combined with --enable-partitioned-dml",
 		},
+		{
+			name: "partitioned_dml_rejects_non_dml",
+			o:    opts{EnablePartitionedDML: true},
+			mode: single{spanner.StrongRead()},
+			err:  "--enable-partitioned-dml can only be used with DML statements",
+		},
+		{
+			name: "jq_lazy_allows_read_write_dml",
+			o:    opts{JqInputMode: "lazy"},
+			mode: readWrite{},
+		},
+		{
+			name: "partitioned_dml_allows_eager",
+			o:    opts{EnablePartitionedDML: true, JqInputMode: "eager"},
+			mode: partitionedDML{},
+		},
+		{
+			name: "partitioned_dml_rejects_lazy",
+			o:    opts{EnablePartitionedDML: true, JqInputMode: "lazy"},
+			mode: partitionedDML{},
+			err:  "--jq-input-mode=lazy is not supported for partitioned DML",
+		},
 	}
 
 	for _, tt := range tests {
@@ -303,6 +325,20 @@ func TestQueryModeForQuery(t *testing.T) {
 		{name: "dml_with_comment", query: "-- c\nUPDATE T SET X=1", partitionedEnabled: false, tb: tb, wantMode: "readWrite"},
 		{name: "partitioned_dml", query: "UPDATE T SET X=1", partitionedEnabled: true, tb: tb, wantMode: "partitionedDML"},
 		{name: "normal_query", query: "SELECT 1", partitionedEnabled: false, tb: candidateTimestamp, wantMode: "single"},
+		{
+			name:               "partitioned_flag_with_select",
+			query:              "SELECT 1",
+			partitionedEnabled: true,
+			tb:                 candidateTimestamp,
+			wantMode:           "single",
+		},
+		{
+			name:               "partitioned_flag_with_ddl",
+			query:              "CREATE TABLE T (K INT64) PRIMARY KEY (K)",
+			partitionedEnabled: true,
+			tb:                 tb,
+			wantMode:           "single",
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/main.go b/main.go
index 18e46ed..51912c3 100644
--- a/main.go
+++ b/main.go
@@ -185,10 +185,10 @@ func isReadWriteStatement(query string) bool {
 }
 
 func queryModeForQuery(query string, enablePartitionedDML bool, tb spanner.TimestampBound) queryMode {
-	if enablePartitionedDML {
-		return partitionedDML{}
-	}
 	if isReadWriteStatement(query) {
+		if enablePartitionedDML {
+			return partitionedDML{}
+		}
 		return readWrite{}
 	}
 	return single{tb}
@@ -215,6 +215,14 @@ func validateExecutionOptions(o opts, mode queryMode) error {
 			return fmt.Errorf("%s cannot be used with DML statements", flagName)
 		}
 	}
+	if o.EnablePartitionedDML {
+		if _, ok := mode.(partitionedDML); !ok {
+			return fmt.Errorf("--enable-partitioned-dml can only be used with DML statements")
+		}
+	}
+	if _, ok := mode.(partitionedDML); ok && o.JqInputMode == "lazy" {
+		return fmt.Errorf("--jq-input-mode=lazy is not supported for partitioned DML")
+	}
 	return nil
 }