bloated ai stuff but will fix

Author: qykr

internal/compiler/parse_test.go

@@ -84,6 +84,18 @@ const lowerSwitchedOrderQuery = `-- name: LowerSwitchedOrder :many
 SELECT bar FROM foo WHERE bar = $1 AND bat = LOWER($2);
 `
 
+const orderByBindsSchema = `
+CREATE TABLE authors (
+	id BIGSERIAL PRIMARY KEY,
+	name text NOT NULL,
+	bio text
+);
+`
+
+const orderByBindsQuery = `-- name: ListAuthorsColumnSortFnWtihArg :many
+SELECT * FROM authors ORDER BY MOD(id, $1);
+`
+
 type stubAnalyzer struct {
 	analyze func(context.Context, ast.Node, string, []string, *named.ParamSet) (*analysispb.Analysis, error)
 }
@@ -216,6 +228,36 @@ func newLowerSwitchedOrderCompiler(t *testing.T) (*Compiler, *ast.RawStmt) {
 	}, stmts[0].Raw
 }
 
+func newOrderByBindsCompiler(t *testing.T) (*Compiler, *ast.RawStmt) {
+	t.Helper()
+
+	parser := postgresql.NewParser()
+	catalog := postgresql.NewCatalog()
+
+	schema, err := parser.Parse(strings.NewReader(orderByBindsSchema))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := catalog.Build(schema); err != nil {
+		t.Fatal(err)
+	}
+
+	stmts, err := parser.Parse(strings.NewReader(orderByBindsQuery))
+	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()
 
@@ -325,6 +367,28 @@ func assertLowerSwitchedOrderParams(t *testing.T, params []Parameter) {
 	}
 }
 
+func assertOrderByBindsParams(t *testing.T, params []Parameter) {
+	t.Helper()
+
+	if len(params) != 1 {
+		t.Fatalf("expected 1 param, got %d", len(params))
+	}
+
+	param := params[0]
+	if param.Number != 1 {
+		t.Fatalf("param number mismatch: got %d want %d", param.Number, 1)
+	}
+	if param.Column == nil {
+		t.Fatal("param column is nil")
+	}
+	if param.Column.Name != "mod" {
+		t.Fatalf("param name mismatch: got %q want %q", param.Column.Name, "mod")
+	}
+	if param.Column.DataType != "bigint" && param.Column.DataType != "pg_catalog.int8" {
+		t.Fatalf("param type mismatch: got %q want %q or %q", param.Column.DataType, "bigint", "pg_catalog.int8")
+	}
+}
+
 func TestInferQueryPreservesInsertSelectParamNamesWithCTEAndMixedParams(t *testing.T) {
 	t.Parallel()
 
@@ -451,3 +515,15 @@ func TestParseQueryPreservesLowerSwitchedOrderParamTypes(t *testing.T) {
 
 	assertLowerSwitchedOrderParams(t, query.Params)
 }
+
+func TestParseQueryPreservesOrderByBindsModParamType(t *testing.T) {
+	t.Parallel()
+
+	comp, raw := newOrderByBindsCompiler(t)
+	query, err := comp.parseQuery(raw, orderByBindsQuery, opts.Parser{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertOrderByBindsParams(t, query.Params)
+}

internal/compiler/resolve.go

@@ -401,6 +401,164 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 		})
 	}
 
+	resolveColumnType := func(node *ast.ColumnRef) *ast.TypeName {
+		if node == nil {
+			return nil
+		}
+
+		items := stringSlice(node.Fields)
+		var schema, key, alias string
+		switch len(items) {
+		case 1:
+			key = items[0]
+		case 2:
+			alias = items[0]
+			key = items[1]
+		case 3:
+			schema = items[0]
+			alias = items[1]
+			key = items[2]
+		default:
+			return nil
+		}
+
+		search := tables
+		if alias != "" {
+			if original, ok := aliasMap[alias]; ok {
+				search = []*ast.TableName{original}
+			} else {
+				var matches []*ast.TableName
+				for _, fqn := range tables {
+					if fqn.Name == alias {
+						matches = append(matches, fqn)
+					}
+				}
+				if len(matches) == 0 {
+					return nil
+				}
+				search = matches
+			}
+		}
+
+		var found *catalog.Column
+		var foundCount int
+		for _, table := range search {
+			tableSchema := table.Schema
+			if tableSchema == "" {
+				tableSchema = c.DefaultSchema
+			}
+			if schema != "" && tableSchema != schema {
+				continue
+			}
+			tablesBySchema, ok := typeMap[tableSchema]
+			if !ok {
+				continue
+			}
+			colsByTable, ok := tablesBySchema[table.Name]
+			if !ok {
+				continue
+			}
+			col, ok := colsByTable[key]
+			if !ok {
+				continue
+			}
+			found = col
+			foundCount++
+		}
+		if foundCount != 1 || found == nil {
+			return nil
+		}
+
+		typ := found.Type
+		return &typ
+	}
+
+	normalizedFuncArgTypeKey := func(t *ast.TypeName) string {
+		if t == nil {
+			return ""
+		}
+
+		name := strings.ToLower(comp.parser.TypeName(t.Schema, t.Name))
+		switch name {
+		case "bigserial", "serial8":
+			name = "bigint"
+		case "serial", "serial4":
+			name = "integer"
+		case "smallserial", "serial2":
+			name = "smallint"
+		}
+
+		schema := t.Schema
+		if schema == "pg_catalog" {
+			schema = ""
+		}
+		return schema + "|" + name + "|" + strconv.Itoa(arrayDims(t))
+	}
+
+	sameFuncArgType := func(a, b *ast.TypeName) bool {
+		if a == nil || b == nil {
+			return a == nil && b == nil
+		}
+		return normalizedFuncArgTypeKey(a) == normalizedFuncArgTypeKey(b)
+	}
+
+	var resolveFuncExprType func(node ast.Node, targetParamNumber int) *ast.TypeName
+	resolveFuncExprType = func(node ast.Node, targetParamNumber int) *ast.TypeName {
+		switch n := node.(type) {
+		case *ast.NamedArgExpr:
+			return resolveFuncExprType(n.Arg, targetParamNumber)
+		case *ast.TypeCast:
+			if pr, ok := n.Arg.(*ast.ParamRef); ok && pr.Number == targetParamNumber {
+				return nil
+			}
+			return n.TypeName
+		case *ast.ColumnRef:
+			return resolveColumnType(n)
+		default:
+			return nil
+		}
+	}
+
+	filterFuncCallCandidates := func(call *ast.FuncCall, funcs []catalog.Function, targetIdx int, targetNamedArg string, targetParamNumber int) []catalog.Function {
+		if call == nil || call.Args == nil || len(funcs) <= 1 {
+			return funcs
+		}
+
+		filtered := funcs
+		for i, item := range call.Args.Items {
+			if i == targetIdx {
+				continue
+			}
+
+			namedArg := ""
+			if named, ok := item.(*ast.NamedArgExpr); ok && named.Name != nil {
+				namedArg = *named.Name
+			}
+			if targetNamedArg != "" && namedArg == targetNamedArg {
+				continue
+			}
+
+			typ := resolveFuncExprType(item, targetParamNumber)
+			if typ == nil {
+				continue
+			}
+
+			narrowed := make([]catalog.Function, 0, len(filtered))
+			for _, fn := range filtered {
+				arg := funcCallArg(fn, i, namedArg)
+				if arg == nil || !sameFuncArgType(arg.Type, typ) {
+					continue
+				}
+				narrowed = append(narrowed, fn)
+			}
+			if len(narrowed) > 0 {
+				filtered = narrowed
+			}
+		}
+
+		return filtered
+	}
+
 	for _, ref := range args {
 		switch n := ref.parent.(type) {
 
@@ -697,7 +855,8 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 					continue
 				}
 
-				paramName, paramType := funcCallArgMetadata(funcs, i, argName)
+				candidateFuncs := filterFuncCallCandidates(n, funcs, i, argName, ref.ref.Number)
+				paramName, paramType := funcCallArgMetadata(candidateFuncs, i, argName)
 				if argName != "" {
 					paramName = argName
 				}