fix: attribution fixes from production corpuses (up to 94% success rate)

Heavily done with use of Claude Code, to reduce the testing times

>   Root causes fixed:
>
>  - Walker node-type gaps: rewriteRefs, walkFixups, walkCast each
>    missed CoalesceExpr / MinMaxExpr / AArrayExpr / RowExpr / NullTest /
>    BooleanTest, plus InsertStmt.WithClause in every walker. Refs and
>    EXTRACT fields buried in COALESCE survived reshape; CTEs on INSERT
>    were invisible to the whole chain.
>  - Transaction cascade: a mid-batch error in attributeCluster's
>    BEGIN; …; COMMIT; left the connection aborted and poisoned every
>    subsequent cluster with 25P02. ROLLBACK on error.
>  - Confidence label: unqualified column on a scan node was labeled
>    heuristic, but PG's plan structure pins it to that relation
>    exactly. Only mismatched qualifiers are heuristic now.
>  - Function Scan: use Alias as table when Relation Name is absent
>    (pg_catalog.pg_settings et al).
>  - Capture: skip session/meta statements so qshape's own probes
>    don't re-enter the corpus.
>
>  Typecast additions from real corpus errors: BoolExpr→bool,
>  json_build_object keys→text, EXISTS SELECT→int4, JOIN ON / WHERE /
>  HAVING→bool, searched CASE WHEN→bool.

Author: radim

cmd/qshape/attribute.go

@@ -100,6 +100,7 @@ func runAttribute(inPath, connStr string, top int) error {
 	}
 	defer conn.Close(ctx)
 
+	cache := newTypecastCache(conn)
 	attributed, skipped := 0, 0
 	for i := range doc.Clusters {
 		if top > 0 && i >= top {
@@ -109,7 +110,7 @@ func runAttribute(inPath, connStr string, top int) error {
 		if c.Fingerprint == "" || c.Canonical == "" {
 			continue
 		}
-		params, err := attributeCluster(ctx, conn, c.Canonical)
+		params, err := attributeCluster(ctx, conn, cache, c.Canonical)
 		if err != nil {
 			skipped++
 			c.Params = []qshape.ParamAttribution{{Confidence: "none", Note: err.Error()}}
@@ -129,10 +130,45 @@ func runAttribute(inPath, connStr string, top int) error {
 	return enc.Encode(doc)
 }
 
-func attributeCluster(ctx context.Context, conn *pgx.Conn, canonical string) ([]qshape.ParamAttribution, error) {
-	var planJSON []byte
-	row := conn.QueryRow(ctx, "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+canonical)
-	if err := row.Scan(&planJSON); err != nil {
+func attributeCluster(ctx context.Context, conn *pgx.Conn, cache *typecastCache, canonical string) ([]qshape.ParamAttribution, error) {
+	// Re-normalise so clusters.json written by an older qshape version picks
+	// up current reshape fixes (extract-field recovery, param renumbering).
+	// Fall back to the stored form if parsing fails
+	if renormed, err := qshape.Normalize(canonical); err == nil {
+		canonical = renormed
+	}
+	explainSQL := castFuncParamRefs(ctx, cache, canonical)
+	// PREPARE + EXPLAIN EXECUTE so Postgres sets up a parameter context
+	// for the $N placeholders. Works on any PG version (GENERIC_PLAN alone
+	// requires 16+, and the simple-query parser rejects bare $N otherwise).
+	// Types come from the typecast pass we just ran; NULL values satisfy
+	// EXECUTE's arity requirement without affecting the plan
+	nparams := maxParamNumber(explainSQL)
+	nulls := "NULL"
+	for i := 1; i < nparams; i++ {
+		nulls += ", NULL"
+	}
+	// force_generic_plan keeps $N in the plan output instead of inlining
+	// the NULL arguments. Without it Postgres produces a custom plan with
+	// `WHERE col = NULL` filters that walkPlan can't attribute
+	script := "SET LOCAL plan_cache_mode = force_generic_plan;\n"
+	script += "PREPARE _qshape_tmp AS " + explainSQL + ";\n"
+	if nparams > 0 {
+		script += "EXPLAIN (FORMAT JSON) EXECUTE _qshape_tmp(" + nulls + ");\n"
+	} else {
+		script += "EXPLAIN (FORMAT JSON) EXECUTE _qshape_tmp;\n"
+	}
+	script += "DEALLOCATE _qshape_tmp;"
+	// SET LOCAL only applies inside a transaction — wrap the whole script
+	script = "BEGIN;\n" + script + "\nCOMMIT;"
+	planJSON, err := readPlanJSON(ctx, conn, script)
+	if err != nil {
+		// A mid-batch error aborts the BEGIN'd transaction and skips the
+		// trailing COMMIT in the same simple-query batch, so the connection
+		// stays in aborted state and every following cluster fails with
+		// 25P02. ROLLBACK resets it before the next call.
+		_, _ = conn.Exec(ctx, "ROLLBACK")
+		_, _ = conn.Exec(ctx, "DEALLOCATE IF EXISTS _qshape_tmp")
 		return nil, err
 	}
 
@@ -157,6 +193,73 @@ func attributeCluster(ctx context.Context, conn *pgx.Conn, canonical string) ([]
 	return out, nil
 }
 
+// readPlanJSON runs a multi-statement script (PREPARE; EXPLAIN; DEALLOCATE)
+// through simple-query protocol and returns the first text cell of the
+// first result-bearing statement. pgx's Query path sanitises $N against
+// bound args; here we send the script verbatim
+func readPlanJSON(ctx context.Context, conn *pgx.Conn, script string) ([]byte, error) {
+	mrr := conn.PgConn().Exec(ctx, script)
+	defer mrr.Close()
+	var out []byte
+	for mrr.NextResult() {
+		rr := mrr.ResultReader()
+		for rr.NextRow() {
+			vals := rr.Values()
+			if len(vals) > 0 && out == nil {
+				out = append([]byte(nil), vals[0]...)
+			}
+		}
+		if _, err := rr.Close(); err != nil {
+			return nil, err
+		}
+	}
+	if err := mrr.Close(); err != nil {
+		return nil, err
+	}
+	if out == nil {
+		return nil, fmt.Errorf("no rows returned")
+	}
+	return out, nil
+}
+
+// maxParamNumber scans sql for `$N` tokens and returns the highest N seen,
+// or 0 if none. Skips $-tags inside string and dollar-quoted contexts
+func maxParamNumber(sql string) int {
+	max := 0
+	for i := 0; i < len(sql); i++ {
+		c := sql[i]
+		switch c {
+		case '\'':
+			// skip to matching quote, honoring doubled ''
+			i++
+			for i < len(sql) {
+				if sql[i] == '\'' {
+					if i+1 < len(sql) && sql[i+1] == '\'' {
+						i += 2
+						continue
+					}
+					break
+				}
+				i++
+			}
+		case '$':
+			j := i + 1
+			n := 0
+			for j < len(sql) && sql[j] >= '0' && sql[j] <= '9' {
+				n = n*10 + int(sql[j]-'0')
+				j++
+			}
+			if j > i+1 {
+				if n > max {
+					max = n
+				}
+				i = j - 1
+			}
+		}
+	}
+	return max
+}
+
 func walkPlan(raw json.RawMessage, parentSchema, parentTable string, ctx *attrCtx) {
 	if len(raw) == 0 {
 		return
@@ -166,21 +269,32 @@ func walkPlan(raw json.RawMessage, parentSchema, parentTable string, ctx *attrCt
 		return
 	}
 
-	// Track alias → table mapping so we can resolve `u.id = $1` to users.id
+	// Track alias → table mapping so we can resolve `u.id = $1` to users.id.
+	// Function Scan on a system view like pg_catalog.pg_settings leaves
+	// RelationName empty but still sets Alias; use Alias as the table name
+	// so conds like `(name = $1)` attribute to pg_settings.name.
 	aliasToTable := map[string]tableRef{}
+	fallbackTable := n.RelationName
+	fallbackSchema := n.Schema
 	if n.RelationName != "" {
 		t := tableRef{Schema: n.Schema, Table: n.RelationName}
 		aliasToTable[n.RelationName] = t
 		if n.Alias != "" && n.Alias != n.RelationName {
 			aliasToTable[n.Alias] = t
 		}
+	} else if n.Alias != "" {
+		// Function Scan / Values Scan: no Relation Name but Alias names the
+		// logical target (e.g. pg_settings). Attribute to the alias.
+		t := tableRef{Schema: n.Schema, Table: n.Alias}
+		aliasToTable[n.Alias] = t
+		fallbackTable = n.Alias
 	}
 
 	for _, cond := range []string{n.IndexCond, n.HashCond, n.Filter, n.RecheckCond, n.JoinFilter, n.MergeCond} {
 		if cond == "" {
 			continue
 		}
-		attributeCond(cond, aliasToTable, n.Schema, n.RelationName, ctx)
+		attributeCond(cond, aliasToTable, fallbackSchema, fallbackTable, ctx)
 	}
 
 	if len(n.Plans) > 0 {
@@ -223,10 +337,19 @@ func recordParam(aliasOrTable, col, posStr string, aliases map[string]tableRef,
 	ref, ok := aliases[aliasOrTable]
 	confidence := "exact"
 	if !ok {
-		// Bare column without an alias — attribute to the current relation.
 		if fallbackTable != "" {
 			ref = tableRef{Schema: fallbackSchema, Table: fallbackTable}
-			confidence = "heuristic"
+			// An unqualified column in plan text (e.g. `Filter: (id = $1)`
+			// on an Index Scan over `session`) is PG telling us exactly
+			// which scan node the column belongs to — not a guess. Only
+			// downgrade to heuristic when we saw a qualifier that didn't
+			// resolve (stale alias, schema-qualified name we didn't
+			// track, or a subplan reference).
+			if aliasOrTable == "" {
+				confidence = "exact"
+			} else {
+				confidence = "heuristic"
+			}
 		} else {
 			confidence = "none"
 		}

cmd/qshape/attribute_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"testing"
 
 	"github.com/boringsql/qshape"
@@ -25,20 +26,42 @@ func TestAttributeCondAliasedEqual(t *testing.T) {
 	}
 }
 
-func TestAttributeCondUnaliasedFallback(t *testing.T) {
+// PG emits bare column names in plan text when the scan is unambiguous
+// (e.g. Filter on a single-table Index Scan). The plan node pins the
+// column to its relation — that's exact attribution, not a guess.
+func TestAttributeCondUnqualifiedOnScanIsExact(t *testing.T) {
 	ctx := &attrCtx{byPosition: map[int]*qshape.ParamAttribution{}}
-	aliases := map[string]tableRef{}
+	aliases := map[string]tableRef{
+		"session": {Schema: "auth", Table: "session"},
+	}
 	attributeCond("(id = $1)", aliases, "auth", "session", ctx)
 
 	a, ok := ctx.byPosition[1]
 	if !ok {
 		t.Fatal("expected param 1 attributed")
 	}
-	if a.Table != "session" || a.Column != "id" {
-		t.Errorf("wrong fallback attribution: %+v", a)
+	if a.Table != "session" || a.Column != "id" || a.Schema != "auth" {
+		t.Errorf("wrong attribution: %+v", a)
+	}
+	if a.Confidence != "exact" {
+		t.Errorf("expected exact confidence for unqualified col on a scan node, got %s", a.Confidence)
+	}
+}
+
+// A qualifier that doesn't resolve (outer-scope ref, schema-qualified name,
+// subplan name) — attribute to the current relation as a best guess and
+// flag it as heuristic.
+func TestAttributeCondMismatchedQualifierIsHeuristic(t *testing.T) {
+	ctx := &attrCtx{byPosition: map[int]*qshape.ParamAttribution{}}
+	aliases := map[string]tableRef{}
+	attributeCond("(outer_alias.id = $1)", aliases, "auth", "session", ctx)
+
+	a, ok := ctx.byPosition[1]
+	if !ok {
+		t.Fatal("expected param 1 attributed")
 	}
 	if a.Confidence != "heuristic" {
-		t.Errorf("expected heuristic confidence, got %s", a.Confidence)
+		t.Errorf("expected heuristic for mismatched qualifier, got %s", a.Confidence)
 	}
 }
 
@@ -59,6 +82,32 @@ func TestAttributeCondMultipleParams(t *testing.T) {
 	}
 }
 
+// PG plans system views like pg_catalog.pg_settings as a Function Scan.
+// The plan node has no "Relation Name" but carries the view name in
+// "Alias". walkPlan must still attribute conds on that node.
+func TestWalkPlanFunctionScanWithAliasOnly(t *testing.T) {
+	// Index Scan on a function-backed view: Relation Name absent, Alias set.
+	plan := json.RawMessage(`{
+		"Node Type": "Function Scan",
+		"Function Name": "pg_show_all_settings",
+		"Alias": "pg_settings",
+		"Filter": "(name = $1)"
+	}`)
+	c := &attrCtx{byPosition: map[int]*qshape.ParamAttribution{}}
+	walkPlan(plan, "", "", c)
+
+	a, ok := c.byPosition[1]
+	if !ok {
+		t.Fatal("expected param 1 attributed via Alias fallback")
+	}
+	if a.Table != "pg_settings" || a.Column != "name" {
+		t.Errorf("wrong attribution: %+v", a)
+	}
+	if a.Confidence != "exact" {
+		t.Errorf("expected exact, got %s", a.Confidence)
+	}
+}
+
 func TestAttributeCondPreservesExactOverHeuristic(t *testing.T) {
 	ctx := &attrCtx{byPosition: map[int]*qshape.ParamAttribution{}}
 	aliases := map[string]tableRef{

cmd/qshape/capture.go

@@ -63,6 +63,10 @@ FROM pg_stat_statements s`)
 	}
 	where = append(where, fmt.Sprintf("s.calls > $%d", len(args)+1))
 	args = append(args, minCalls)
+	// Skip session/meta statements: they can't be PREPAREd (DISCARD,
+	// BEGIN/COMMIT, SET, SHOW, VACUUM, ANALYZE, reload) and qshape's own
+	// EXPLAIN / PREPARE / DEALLOCATE probes pollute pg_stat_statements.
+	where = append(where, `s.query !~* '^\s*(discard|begin|start|commit|rollback|savepoint|release|set|reset|show|listen|unlisten|notify|checkpoint|vacuum|analyze|reindex|cluster|explain|prepare|deallocate|execute|close|fetch|move|lock)\M'`)
 	sb.WriteString("\nWHERE ")
 	sb.WriteString(strings.Join(where, " AND "))
 	sb.WriteString("\nORDER BY s.total_exec_time DESC")