From 18104e02e2fd3985fcb6498852b6ffe408917383 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 13:13:46 -0400 Subject: [PATCH 1/4] fix flight sql explain plan rows --- duckdbservice/arrow_helpers.go | 43 ++++++++++++++++++++++--- duckdbservice/arrow_helpers_test.go | 50 ++++++++++++++++++++++++++++- 2 files changed, 87 insertions(+), 6 deletions(-) diff --git a/duckdbservice/arrow_helpers.go b/duckdbservice/arrow_helpers.go index 817d1523..083ab049 100644 --- a/duckdbservice/arrow_helpers.go +++ b/duckdbservice/arrow_helpers.go @@ -36,7 +36,10 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema, builder := array.NewRecordBuilder(alloc, schema) defer builder.Release() - numFields := schema.NumFields() + scanFields, explainValueColumn, err := rowScanShape(rows, schema) + if err != nil { + return nil, err + } count := 0 // Order matters: check `count < batchSize` first, then call rows.Next(). // The reverse (rows.Next() && count < batchSize) advances the cursor once @@ -47,8 +50,8 @@ 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) + values := make([]interface{}, scanFields) + valuePtrs := make([]interface{}, scanFields) for i := range values { valuePtrs[i] = &values[i] } @@ -57,8 +60,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 +79,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) } diff --git a/duckdbservice/arrow_helpers_test.go b/duckdbservice/arrow_helpers_test.go index 1234afb9..2f826a63 100644 --- a/duckdbservice/arrow_helpers_test.go +++ b/duckdbservice/arrow_helpers_test.go @@ -1069,6 +1069,55 @@ 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() + query := "EXPLAIN SELECT 1" + schema, err := GetQuerySchema(ctx, conn, 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 != "physical_plan" { + t.Fatalf("schema field = %q, want physical_plan", got) + } + + rows, err := conn.QueryContext(ctx, 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() == 0 { + t.Fatal("EXPLAIN returned no plan rows") + } + if got := plans.Value(0); got == "" || got == "physical_plan" { + t.Fatalf("first plan row = %q, want rendered plan text", got) + } +} + 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 +1255,3 @@ func TestRowsToRecordNoRowsLostAtBatchBoundary(t *testing.T) { }) } } - From b723ff6779d5c593752f1803d5cab0f3d42a7641 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 13:19:07 -0400 Subject: [PATCH 2/4] test explain analyze plan row conversion --- duckdbservice/arrow_helpers_test.go | 107 +++++++++++++++++++++------- 1 file changed, 81 insertions(+), 26 deletions(-) diff --git a/duckdbservice/arrow_helpers_test.go b/duckdbservice/arrow_helpers_test.go index 2f826a63..0ddcb3d8 100644 --- a/duckdbservice/arrow_helpers_test.go +++ b/duckdbservice/arrow_helpers_test.go @@ -1082,39 +1082,94 @@ func TestRowsToRecordExplainUsesPlanValueColumn(t *testing.T) { defer func() { _ = conn.Close() }() ctx := context.Background() - query := "EXPLAIN SELECT 1" - schema, err := GetQuerySchema(ctx, conn, 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 _, err := conn.ExecContext(ctx, "PRAGMA explain_output='all'"); err != nil { + t.Fatalf("set explain output: %v", err) } - if got := schema.Field(0).Name; got != "physical_plan" { - t.Fatalf("schema field = %q, want physical_plan", got) + if _, err := conn.ExecContext(ctx, "CREATE TABLE explain_write_once (id INTEGER)"); err != nil { + t.Fatalf("create write table: %v", err) } - rows, err := conn.QueryContext(ctx, query) - if err != nil { - t.Fatalf("query: %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) + } + }, + }, } - 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() + 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() }() - plans := rec.Column(0).(*array.String) - if plans.Len() == 0 { - t.Fatal("EXPLAIN returned no plan rows") + 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) + } + } + if tc.after != nil { + tc.after(t) + } + }) } - if got := plans.Value(0); got == "" || got == "physical_plan" { - t.Fatalf("first plan row = %q, want rendered plan text", got) +} + +func isExplainKey(value string) bool { + switch value { + case "logical_plan", "logical_opt", "physical_plan", "analyzed_plan": + return true + default: + return false } } From f1aea2c15ff5ecbe3bc249040134343872ce1ace Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 13:46:29 -0400 Subject: [PATCH 3/4] test: add e2e DuckLake EXPLAIN regression --- tests/e2e-mw-dev/README.md | 3 ++- tests/e2e-mw-dev/harness.sh | 21 +++++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/e2e-mw-dev/README.md b/tests/e2e-mw-dev/README.md index 346589c4..05e25309 100644 --- a/tests/e2e-mw-dev/README.md +++ b/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` diff --git a/tests/e2e-mw-dev/harness.sh b/tests/e2e-mw-dev/harness.sh index 181ac1a7..65ddc33d 100755 --- a/tests/e2e-mw-dev/harness.sh +++ b/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) @@ -736,6 +736,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 @@ -2273,6 +2284,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" @@ -2345,8 +2357,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. @@ -2480,7 +2493,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 + 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 + isolation + lifecycle-teardown, on cnpg & ext (4 parallel lanes)" } main "$@" From 19bf7b628e9c84c27bfb7f0de23b67060b818a4b Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 15:09:12 -0400 Subject: [PATCH 4/4] fix explain rows eof conversion --- duckdbservice/arrow_helpers.go | 18 ++++++++++++++---- duckdbservice/arrow_helpers_test.go | 9 +++++++++ 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/duckdbservice/arrow_helpers.go b/duckdbservice/arrow_helpers.go index 083ab049..ab66e854 100644 --- a/duckdbservice/arrow_helpers.go +++ b/duckdbservice/arrow_helpers.go @@ -36,10 +36,9 @@ func RowsToRecord(alloc memory.Allocator, rows *sql.Rows, schema *arrow.Schema, builder := array.NewRecordBuilder(alloc, schema) defer builder.Release() - scanFields, explainValueColumn, err := rowScanShape(rows, schema) - if err != nil { - return nil, err - } + 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 @@ -50,6 +49,17 @@ 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() { + // 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 { diff --git a/duckdbservice/arrow_helpers_test.go b/duckdbservice/arrow_helpers_test.go index 0ddcb3d8..cf4013b2 100644 --- a/duckdbservice/arrow_helpers_test.go +++ b/duckdbservice/arrow_helpers_test.go @@ -1157,6 +1157,15 @@ func TestRowsToRecordExplainUsesPlanValueColumn(t *testing.T) { 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) }