[codex] fix flight sql explain plan rows

duckdbservice/arrow_helpers.go

@@ -36,7 +36,9 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema,
 	builder := array.NewRecordBuilder(alloc, schema)
 	defer builder.Release()
 
-	numFields := schema.NumFields()
+	scanFields := schema.NumFields()
+	explainValueColumn := -1
+	scanShapeResolved := false
 	count := 0
 	// Order matters: check `count < batchSize` first, then call rows.Next().
 	// The reverse (rows.Next() && count < batchSize) advances the cursor once
@@ -47,8 +49,19 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema,
 	// parquet-metadata row count, so the discrepancy was invisible to
 	// aggregation queries. See TestRowsToRecordNoRowsLostAtBatchBoundary.
 	for count < batchSize && rows.Next() {
-		values := make([]interface{}, numFields)
-		valuePtrs := make([]interface{}, numFields)
+		// Resolve the physical scan shape only after Next succeeds. Some drivers
+		// close rows automatically at EOF; calling Columns after that would turn
+		// normal completion into "sql: Rows are closed".
+		if !scanShapeResolved {
+			var err error
+			scanFields, explainValueColumn, err = rowScanShape(rows, schema)
+			if err != nil {
+				return nil, err
+			}
+			scanShapeResolved = true
+		}
+		values := make([]interface{}, scanFields)
+		valuePtrs := make([]interface{}, scanFields)
 		for i := range values {
 			valuePtrs[i] = &values[i]
 		}
@@ -57,8 +70,12 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema,
 			return nil, err
 		}
 
-		for i, val := range values {
-			AppendValue(builder.Field(i), val)
+		if explainValueColumn >= 0 {
+			AppendValue(builder.Field(0), values[explainValueColumn])
+		} else {
+			for i, val := range values {
+				AppendValue(builder.Field(i), val)
+			}
 		}
 		count++
 	}
@@ -72,6 +89,32 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema,
 	return builder.NewRecordBatch(), nil
 }
 
+func rowScanShape(rows *sql.Rows, schema *arrow.Schema) (scanFields int, explainValueColumn int, err error) {
+	scanFields = schema.NumFields()
+	explainValueColumn = -1
+	if !isExplainPlanSchema(schema) {
+		return scanFields, explainValueColumn, nil
+	}
+	cols, err := rows.Columns()
+	if err != nil {
+		return 0, -1, err
+	}
+	if len(cols) == 2 &&
+		strings.EqualFold(cols[0], "explain_key") &&
+		strings.EqualFold(cols[1], "explain_value") {
+		return 2, 1, nil
+	}
+	return scanFields, explainValueColumn, nil
+}
+
+func isExplainPlanSchema(schema *arrow.Schema) bool {
+	if schema == nil || schema.NumFields() != 1 {
+		return false
+	}
+	name := schema.Field(0).Name
+	return name == "physical_plan" || name == "analyzed_plan"
+}
+
 type contextQueryer interface {
 	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
 }

duckdbservice/arrow_helpers_test.go

@@ -1069,6 +1069,119 @@ func TestGetQuerySchemaExplainDoesNotExecute(t *testing.T) {
 	}
 }
 
+func TestRowsToRecordExplainUsesPlanValueColumn(t *testing.T) {
+	db, err := sql.Open("duckdb", "")
+	if err != nil {
+		t.Fatalf("failed to open DuckDB: %v", err)
+	}
+	defer func() { _ = db.Close() }()
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		t.Fatalf("open conn: %v", err)
+	}
+	defer func() { _ = conn.Close() }()
+
+	ctx := context.Background()
+	if _, err := conn.ExecContext(ctx, "PRAGMA explain_output='all'"); err != nil {
+		t.Fatalf("set explain output: %v", err)
+	}
+	if _, err := conn.ExecContext(ctx, "CREATE TABLE explain_write_once (id INTEGER)"); err != nil {
+		t.Fatalf("create write table: %v", err)
+	}
+
+	cases := []struct {
+		name        string
+		query       string
+		wantCol     string
+		minPlanRows int
+		after       func(t *testing.T)
+	}{
+		{
+			name:        "explain",
+			query:       "EXPLAIN SELECT 1",
+			wantCol:     "physical_plan",
+			minPlanRows: 2,
+		},
+		{
+			name:        "explain analyze",
+			query:       "EXPLAIN ANALYZE INSERT INTO explain_write_once VALUES (1)",
+			wantCol:     "analyzed_plan",
+			minPlanRows: 1,
+			after: func(t *testing.T) {
+				t.Helper()
+				var n int
+				if err := conn.QueryRowContext(ctx, "SELECT count(*) FROM explain_write_once").Scan(&n); err != nil {
+					t.Fatalf("count write table: %v", err)
+				}
+				if n != 1 {
+					t.Fatalf("EXPLAIN ANALYZE write count = %d, want 1", n)
+				}
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			schema, err := GetQuerySchema(ctx, conn, tc.query, nil)
+			if err != nil {
+				t.Fatalf("GetQuerySchema: %v", err)
+			}
+			if got := schema.NumFields(); got != 1 {
+				t.Fatalf("schema fields = %d, want 1", got)
+			}
+			if got := schema.Field(0).Name; got != tc.wantCol {
+				t.Fatalf("schema field = %q, want %s", got, tc.wantCol)
+			}
+
+			rows, err := conn.QueryContext(ctx, tc.query)
+			if err != nil {
+				t.Fatalf("query: %v", err)
+			}
+			defer func() { _ = rows.Close() }()
+
+			rec, err := RowsToRecord(memory.NewGoAllocator(), rows, schema, 1024)
+			if err != nil {
+				t.Fatalf("RowsToRecord: %v", err)
+			}
+			if rec == nil {
+				t.Fatal("RowsToRecord returned nil")
+			}
+			defer rec.Release()
+
+			plans := rec.Column(0).(*array.String)
+			if plans.Len() < tc.minPlanRows {
+				t.Fatalf("EXPLAIN returned %d plan rows, want at least %d", plans.Len(), tc.minPlanRows)
+			}
+			for i := 0; i < plans.Len(); i++ {
+				if got := plans.Value(i); got == "" || isExplainKey(got) {
+					t.Fatalf("plan row %d = %q, want rendered plan text", i, got)
+				}
+			}
+
+			next, err := RowsToRecord(memory.NewGoAllocator(), rows, schema, 1024)
+			if err != nil {
+				t.Fatalf("RowsToRecord at EOF: %v", err)
+			}
+			if next != nil {
+				defer next.Release()
+				t.Fatal("RowsToRecord at EOF returned a record, want nil")
+			}
+			if tc.after != nil {
+				tc.after(t)
+			}
+		})
+	}
+}
+
+func isExplainKey(value string) bool {
+	switch value {
+	case "logical_plan", "logical_opt", "physical_plan", "analyzed_plan":
+		return true
+	default:
+		return false
+	}
+}
+
 func TestGetQuerySchemaTrailingSemicolon(t *testing.T) {
 	// Regression test: queries ending with ";" caused "syntax error at or near LIMIT"
 	// because GetQuerySchema appended " LIMIT 0" after the semicolon, producing "; LIMIT 0".
@@ -1206,4 +1319,3 @@ func TestRowsToRecordNoRowsLostAtBatchBoundary(t *testing.T) {
 		})
 	}
 }
-

tests/e2e-mw-dev/README.md

@@ -63,7 +63,8 @@ client-go:
   / `timed out waiting for an available worker`) rather than a hang/500/drop, and
   the pool must then serve a retrying connection. The harness logs whether
   backpressure was observed **and** handles it (queries retry through it).
-- **activation** — DuckLake catalogs attach and read/write on cnpg and ext.
+- **activation** — DuckLake catalogs attach, read/write, and run
+  `EXPLAIN`/`EXPLAIN ANALYZE` on cnpg and ext.
 - **worker sizing** (TTL-pool model, `docs/design/worker-ttl-pool.md`) — a
   client-sized connection (`duckgres.worker_cpu`/`worker_memory`/`worker_ttl`
   startup options, sent via `PGOPTIONS`; CP runs `allowClientWorkerProfile=true`

tests/e2e-mw-dev/harness.sh

@@ -26,7 +26,7 @@
 #                  transpilation (#716), and a pipelined extended-query error
 #                  discards queued messages until Sync (#718). A same pgwire
 #                  session remains usable immediately after CancelRequest.
-#   activation   : DuckLake catalogs attach and read/write on cnpg + ext.
+#   activation   : DuckLake catalogs attach, read/write, and EXPLAIN/ANALYZE on cnpg + ext.
 #   sizing       : a client-sized connection (duckgres.worker_cpu/memory/ttl)
 #                  spawns a worker pod carrying the requested CPU+memory, and a
 #                  same-shape reconnect reuses that hot-idle worker (no respawn)
@@ -771,6 +771,17 @@ rw_ducklake() { # org password
   pg "$1" "$2" ducklake "DROP TABLE $t;"
 }
 
+explain_ducklake() { # org password
+  log "EXPLAIN on DuckLake table on $1"
+  t="e2e_explain_$(echo "$1" | tr -c 'a-z0-9' _)"
+  pg "$1" "$2" ducklake "DROP TABLE IF EXISTS $t; CREATE TABLE $t AS SELECT i AS id FROM generate_series(1,3) t(i);"
+  out="$(pg "$1" "$2" ducklake "EXPLAIN SELECT count(*) FROM $t WHERE id > 1;")"
+  [ -n "$out" ] || fail "$1 EXPLAIN returned empty output"
+  out="$(pg "$1" "$2" ducklake "EXPLAIN ANALYZE SELECT count(*) FROM $t WHERE id > 1;")"
+  [ -n "$out" ] || fail "$1 EXPLAIN ANALYZE returned empty output"
+  pg "$1" "$2" ducklake "DROP TABLE $t;"
+}
+
 # Regression for the S3-throttling backfill failure: a sustained HTTP 503
 # SlowDown during a DuckLake DELETE used to outlast httpfs's tiny default retry
 # budget (http_retries=3, ~0.5s cumulative backoff) and surface as a fatal
@@ -2342,6 +2353,7 @@ lane_cnpg() { # full wire/catalog/concurrency/sizing coverage on the cnpg org
   jsonb_concat_semantics "$CNPG" "$cnpg_pw"
   cold_burst_absorption  "$CNPG" "$cnpg_pw"   # early, while this org is mostly cold
   rw_ducklake            "$CNPG" "$cnpg_pw"
+  explain_ducklake       "$CNPG" "$cnpg_pw"
   httpfs_retry_budget    "$CNPG" "$cnpg_pw"   # S3-503 retry budget raised per worker (applyHTTPFSRetryBudget)
   persistent_user_secret "$CNPG" "$cnpg_pw"   # after rw_ducklake (org worker hot)
   persistent_user_secret_isolation "$CNPG" "$cnpg_pw"
@@ -2414,8 +2426,9 @@ lane_ext() { # external-RDS metadata backend + org default profile
   basic_query  "$EXT" "$ext_pw"
   pg_compat_functions "$EXT" "$ext_pw"
   query_source_guc    "$EXT" "$ext_pw"
-  rw_ducklake  "$EXT" "$ext_pw"
-  httpfs_retry_budget "$EXT" "$ext_pw"      # S3-503 retry budget per worker, ext-metadata backend too
+  rw_ducklake         "$EXT" "$ext_pw"
+  explain_ducklake    "$EXT" "$ext_pw"
+  httpfs_retry_budget "$EXT" "$ext_pw"    # S3-503 retry budget per worker, ext-metadata backend too
   persistent_user_secret "$EXT" "$ext_pw"   # secret replay on the ext-metadata org too
   # Org default profile on ext: no client-sized assertions run on this org, so
   # the 2-CPU shape is unambiguously the org default's.
@@ -2550,7 +2563,7 @@ main() {
   # mid-run image bump); it stays covered by the controlplane/ unit tests.
   log "SKIP version-reaper (needs an in-run image bump; see README)"
 
-  log "PASS: admin-no-query-token + models-explorer-api(redaction) + admin-console-api(me/live/metrics/auth-gate) + admin-rbac-viewer(403 mutate/audit) + admin-impersonation(round-trip+audit) + wire + malformed-startup-resilience + jsonb-concat + cold-burst-absorption + pipeline-error-recovery + cancel-reuse + activation(DuckLake) + ext-forks + worker-pod + concurrency + durability + crash-recovery + busy-only-do-not-disrupt + graceful-drain + one-session-per-worker + parallel-cold-burst-ramp + worker-sizing(cnpg DuckLake) + org-default-profile(ext) + persistent-user-secrets(cnpg+ext, cross-user isolation) + user-kill-switch(cnpg) + user-disable-block(cnpg+ext) + connection-duration-logged + compute-usage-pull-api(cnpg) + isolation + lifecycle-teardown, on cnpg & ext (4 parallel lanes)"
+  log "PASS: admin-no-query-token + models-explorer-api(redaction) + admin-console-api(me/live/metrics/auth-gate) + admin-rbac-viewer(403 mutate/audit) + admin-impersonation(round-trip+audit) + wire + malformed-startup-resilience + jsonb-concat + cold-burst-absorption + pipeline-error-recovery + cancel-reuse + activation(DuckLake) + ducklake-explain + ext-forks + worker-pod + concurrency + durability + crash-recovery + busy-only-do-not-disrupt + graceful-drain + one-session-per-worker + parallel-cold-burst-ramp + worker-sizing(cnpg DuckLake) + org-default-profile(ext) + persistent-user-secrets(cnpg+ext, cross-user isolation) + user-kill-switch(cnpg) + user-disable-block(cnpg+ext) + connection-duration-logged + compute-usage-pull-api(cnpg) + isolation + lifecycle-teardown, on cnpg & ext (4 parallel lanes)"
 }
 
 main "$@"