diff --git a/duckdbservice/arrow_helpers.go b/duckdbservice/arrow_helpers.go index 817d1523..ab66e854 100644 --- a/duckdbservice/arrow_helpers.go +++ b/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) } diff --git a/duckdbservice/arrow_helpers_test.go b/duckdbservice/arrow_helpers_test.go index 1234afb9..cf4013b2 100644 --- a/duckdbservice/arrow_helpers_test.go +++ b/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) { }) } } - diff --git a/tests/e2e-mw-dev/README.md b/tests/e2e-mw-dev/README.md index 6dfb4dae..d6518303 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 62a764b1..ebe5a976 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) @@ -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 "$@"