fix: inferred column names

Author: qykr

internal/compiler/find_params.go

@@ -140,6 +140,11 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 		p.parent = node
 
 	case *ast.SelectStmt:
+		if n.FromClause != nil && len(n.FromClause.Items) == 1 {
+			if rv, ok := n.FromClause.Items[0].(*ast.RangeVar); ok {
+				p.rangeVar = rv
+			}
+		}
 		if n.LimitCount != nil {
 			p.limitCount = n.LimitCount
 		}

internal/compiler/parse_test.go

@@ -0,0 +1,175 @@
+package compiler
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	analysispb "github.com/sqlc-dev/sqlc/internal/analysis"
+	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/engine/postgresql"
+	"github.com/sqlc-dev/sqlc/internal/opts"
+	"github.com/sqlc-dev/sqlc/internal/sql/ast"
+	"github.com/sqlc-dev/sqlc/internal/sql/named"
+)
+
+const batchParameterTypeSchema = `
+CREATE TABLE public.solar_commcard_mapping (
+	id INT8 NOT NULL,
+	"deviceId" INT8 NOT NULL,
+	version VARCHAR(32) DEFAULT ''::VARCHAR NOT NULL,
+	sn VARCHAR(32) DEFAULT ''::VARCHAR NOT NULL,
+	"createdAt" TIMESTAMPTZ DEFAULT now(),
+	"updatedAt" TIMESTAMPTZ DEFAULT now()
+);
+`
+
+const batchParameterTypeQuery = `-- name: InsertMappping :batchexec
+WITH
+	table1 AS (
+		SELECT version
+		FROM solar_commcard_mapping
+		WHERE "deviceId" = $1
+		ORDER BY "updatedAt" DESC
+		LIMIT 1
+	)
+INSERT INTO solar_commcard_mapping ("deviceId", version, sn, "updatedAt")
+SELECT $1, @version::text, $3, $4
+WHERE NOT EXISTS (
+	SELECT *
+	FROM table1
+	WHERE table1.version = @version::text
+) OR NOT EXISTS (SELECT * FROM table1);
+`
+
+type stubAnalyzer struct {
+	analyze func(context.Context, ast.Node, string, []string, *named.ParamSet) (*analysispb.Analysis, error)
+}
+
+func (s stubAnalyzer) Analyze(ctx context.Context, n ast.Node, q string, schema []string, np *named.ParamSet) (*analysispb.Analysis, error) {
+	return s.analyze(ctx, n, q, schema, np)
+}
+
+func (stubAnalyzer) Close(context.Context) error                              { return nil }
+func (stubAnalyzer) EnsureConn(context.Context, []string) error               { return nil }
+func (stubAnalyzer) GetColumnNames(context.Context, string) ([]string, error) { return nil, nil }
+
+func newBatchParameterTypeCompiler(t *testing.T) (*Compiler, *ast.RawStmt) {
+	t.Helper()
+
+	parser := postgresql.NewParser()
+	catalog := postgresql.NewCatalog()
+
+	schema, err := parser.Parse(strings.NewReader(batchParameterTypeSchema))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := catalog.Build(schema); err != nil {
+		t.Fatal(err)
+	}
+
+	stmts, err := parser.Parse(strings.NewReader(batchParameterTypeQuery))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(stmts) != 1 {
+		t.Fatalf("expected 1 statement, got %d", len(stmts))
+	}
+
+	return &Compiler{
+		conf:     config.SQL{Engine: config.EnginePostgreSQL},
+		parser:   parser,
+		catalog:  catalog,
+		selector: newDefaultSelector(),
+	}, stmts[0].Raw
+}
+
+func assertBatchParameterNames(t *testing.T, params []Parameter) {
+	t.Helper()
+
+	checks := []struct {
+		idx      int
+		number   int
+		name     string
+		original string
+		named    bool
+	}{
+		{idx: 0, number: 1, name: "deviceId", original: "deviceId"},
+		{idx: 1, number: 2, name: "version", original: "version", named: true},
+		{idx: 2, number: 3, name: "sn", original: "sn"},
+		{idx: 3, number: 4, name: "updatedAt", original: "updatedAt"},
+	}
+	if len(params) != len(checks) {
+		t.Fatalf("expected %d params, got %d", len(checks), len(params))
+	}
+
+	for _, check := range checks {
+		param := params[check.idx]
+		if param.Number != check.number {
+			t.Fatalf("param %d number mismatch: got %d want %d", check.idx, param.Number, check.number)
+		}
+		if param.Column == nil {
+			t.Fatalf("param %d column is nil", check.idx)
+		}
+		if param.Column.Name != check.name {
+			t.Fatalf("param %d name mismatch: got %q want %q", check.idx, param.Column.Name, check.name)
+		}
+		if param.Column.OriginalName != check.original {
+			t.Fatalf("param %d original name mismatch: got %q want %q", check.idx, param.Column.OriginalName, check.original)
+		}
+		if param.Column.IsNamedParam != check.named {
+			t.Fatalf("param %d named mismatch: got %v want %v", check.idx, param.Column.IsNamedParam, check.named)
+		}
+		if param.Column.DataType == "" || param.Column.DataType == "any" {
+			t.Fatalf("param %d type was not inferred: %+v", check.idx, param.Column)
+		}
+	}
+}
+
+func TestInferQueryPreservesInsertSelectParamNamesWithCTEAndMixedParams(t *testing.T) {
+	t.Parallel()
+
+	comp, raw := newBatchParameterTypeCompiler(t)
+	anlys, err := comp.inferQuery(raw, batchParameterTypeQuery)
+	if err != nil && !strings.Contains(err.Error(), "parameter $2") {
+		t.Fatalf("unexpected infer error: %v", err)
+	}
+	if anlys == nil {
+		t.Fatal("expected non-nil analysis")
+	}
+	if !strings.Contains(anlys.Query, "$2::text") {
+		t.Fatalf("expected rewritten query to contain $2::text, got %q", anlys.Query)
+	}
+
+	assertBatchParameterNames(t, anlys.Parameters)
+}
+
+func TestParseQueryManagedDBPreservesInferredParamNames(t *testing.T) {
+	t.Parallel()
+
+	comp, raw := newBatchParameterTypeCompiler(t)
+	comp.analyzer = stubAnalyzer{analyze: func(_ context.Context, _ ast.Node, query string, _ []string, np *named.ParamSet) (*analysispb.Analysis, error) {
+		if np == nil {
+			t.Fatal("expected named param set")
+		}
+		if got, ok := np.NameFor(2); !ok || got != "version" {
+			t.Fatalf("expected param 2 to be named version, got %q %v", got, ok)
+		}
+		if !strings.Contains(query, "$2::text") {
+			t.Fatalf("expected analyzer query to contain rewritten named param, got %q", query)
+		}
+		return &analysispb.Analysis{Params: []*analysispb.Parameter{
+			{Number: 1, Column: &analysispb.Column{DataType: "pg_catalog.int8"}},
+			{Number: 2, Column: &analysispb.Column{Name: "version", DataType: "text", IsNamedParam: true}},
+			{Number: 3, Column: &analysispb.Column{DataType: "text"}},
+			{Number: 4, Column: &analysispb.Column{DataType: "pg_catalog.timestamptz"}},
+		}}, nil
+	}}
+
+	query, err := comp.parseQuery(raw, batchParameterTypeQuery, opts.Parser{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertBatchParameterNames(t, query.Params)
+}