core: rewrite sqlc.arg() to native bind parameters in a preprocessing pass above the engines (#4537)

Author: kyleconroy

internal/compiler/analyze.go

@@ -8,7 +8,7 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/named"
-	"github.com/sqlc-dev/sqlc/internal/sql/rewrite"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 	"github.com/sqlc-dev/sqlc/internal/sql/validate"
 )
 
@@ -114,15 +114,15 @@ func combineAnalysis(prev *analysis, a *analyzer.Analysis) *analysis {
 	return prev
 }
 
-func (c *Compiler) analyzeQuery(raw *ast.RawStmt, query string) (*analysis, error) {
-	return c._analyzeQuery(raw, query, true)
+func (c *Compiler) analyzeQuery(raw *ast.RawStmt, query string, pre *preprocess.Statement) (*analysis, error) {
+	return c._analyzeQuery(raw, query, pre, true)
 }
 
-func (c *Compiler) inferQuery(raw *ast.RawStmt, query string) (*analysis, error) {
-	return c._analyzeQuery(raw, query, false)
+func (c *Compiler) inferQuery(raw *ast.RawStmt, query string, pre *preprocess.Statement) (*analysis, error) {
+	return c._analyzeQuery(raw, query, pre, false)
 }
 
-func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool) (*analysis, error) {
+func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, pre *preprocess.Statement, failfast bool) (*analysis, error) {
 	errors := make([]error, 0)
 	check := func(err error) error {
 		if failfast {
@@ -134,12 +134,12 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
 		return nil
 	}
 
-	numbers, dollar, err := validate.ParamRef(raw)
-	if err := check(err); err != nil {
+	dollar := pre.Dollar
+	if err := check(pre.ParamErr); err != nil {
 		return nil, err
 	}
 
-	raw, namedParams, edits := rewrite.NamedParameters(c.conf.Engine, raw, numbers, dollar)
+	namedParams := pre.Params
 
 	var table *ast.TableName
 	switch n := raw.Stmt.(type) {
@@ -158,7 +158,7 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
 		return nil, err
 	}
 
-	if err := check(validate.In(c.catalog, raw)); err != nil {
+	if err := check(validate.Slice(raw, pre.Slices)); err != nil {
 		return nil, err
 	}
 	rvs := rangeVars(raw.Stmt)
@@ -175,7 +175,7 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
 	} else {
 		sort.Slice(refs, func(i, j int) bool { return refs[i].ref.Number < refs[j].ref.Number })
 	}
-	raw, embeds := rewrite.Embeds(raw)
+	embeds := pre.Embeds
 	qc, err := c.buildQueryCatalog(c.catalog, raw.Stmt, embeds)
 	if err := check(err); err != nil {
 		return nil, err
@@ -190,11 +190,10 @@ func (c *Compiler) _analyzeQuery(raw *ast.RawStmt, query string, failfast bool)
 		return nil, err
 	}
 
-	expandEdits, err := c.expand(qc, raw)
+	edits, err := c.expand(qc, raw)
 	if check(err); err != nil {
 		return nil, err
 	}
-	edits = append(edits, expandEdits...)
 	expanded, err := source.Mutate(query, edits)
 	if err != nil {
 		return nil, err

internal/compiler/compile.go

@@ -16,6 +16,7 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/rpc"
 	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlerr"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
 )
@@ -103,19 +104,30 @@ func (c *Compiler) parseQueries(o opts.Parser) (*Result, error) {
 			continue
 		}
 		src := string(blob)
-		stmts, err := c.parser.Parse(strings.NewReader(src))
+
+		// sqlc syntax is not SQL. Rewrite it to native placeholders before the
+		// engine parser ever sees the query, so parsers only handle SQL.
+		pp := preprocess.File(c.conf.Engine, src)
+
+		stmts, err := c.parser.Parse(strings.NewReader(pp.Text))
 		if err != nil {
 			merr.Add(filename, src, 0, err)
 			continue
 		}
 		for _, stmt := range stmts {
-			query, err := c.parseQuery(stmt.Raw, src, o)
+			query, err := c.parseQuery(stmt.Raw, pp, o)
 			if err != nil {
 				var e *sqlerr.Error
 				loc := stmt.Raw.Pos()
 				if errors.As(err, &e) && e.Location != 0 {
 					loc = e.Location
 				}
+				// Locations are reported against the rewritten query; map them
+				// back so errors point at what the user wrote.
+				loc = pp.Origin(loc)
+				if e != nil && e.Location != 0 {
+					e.Location = loc
+				}
 				merr.Add(filename, src, loc, err)
 				// If this rpc unauthenticated error bubbles up, then all future parsing/analysis will fail
 				if errors.Is(err, rpc.ErrUnauthenticated) {
@@ -130,7 +142,7 @@ func (c *Compiler) parseQueries(o opts.Parser) (*Result, error) {
 			queryName := query.Metadata.Name
 			if queryName != "" {
 				if _, exists := set[queryName]; exists {
-					merr.Add(filename, src, stmt.Raw.Pos(), fmt.Errorf("duplicate query name: %s", queryName))
+					merr.Add(filename, src, pp.Origin(stmt.Raw.Pos()), fmt.Errorf("duplicate query name: %s", queryName))
 					continue
 				}
 				set[queryName] = struct{}{}

internal/compiler/expand.go

@@ -166,34 +166,27 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 			}
 		}
 
-		var oldString string
-		var oldFunc func(string) int
-
-		// use the sqlc.embed string instead
-		if embed, ok := qc.embeds.Find(ref); ok {
-			oldString = embed.Orig()
-		} else {
-			oldFunc = func(s string) int {
-				length := 0
-				for i, o := range old {
-					if hasSeparator := i > 0; hasSeparator {
-						length++
-					}
-					if strings.HasPrefix(s[length:], o) {
-						length += len(o)
-					} else if quoted := c.quote(o); strings.HasPrefix(s[length:], quoted) {
-						length += len(quoted)
-					} else {
-						length += len(o)
-					}
+		// An embed was rewritten to "table.*" in the query text, so it is
+		// measured the same way as a star reference the user wrote.
+		oldFunc := func(s string) int {
+			length := 0
+			for i, o := range old {
+				if hasSeparator := i > 0; hasSeparator {
+					length++
+				}
+				if strings.HasPrefix(s[length:], o) {
+					length += len(o)
+				} else if quoted := c.quote(o); strings.HasPrefix(s[length:], quoted) {
+					length += len(quoted)
+				} else {
+					length += len(o)
 				}
-				return length
 			}
+			return length
 		}
 
 		edits = append(edits, source.Edit{
 			Location: res.Location - raw.StmtLocation,
-			Old:      oldString,
 			OldFunc:  oldFunc,
 			New:      strings.Join(cols, ", "),
 		})

internal/compiler/output_columns.go

@@ -258,7 +258,7 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 			if hasStarRef(n) {
 
 				// add a column with a reference to an embedded table
-				if embed, ok := qc.embeds.Find(n); ok {
+				if embed, ok := qc.embeds.Find(n.Location, res.Location); ok {
 					cols = append(cols, &Column{
 						Name:       embed.Table.Name,
 						EmbedTable: embed.Table,

internal/compiler/parse.go

@@ -12,15 +12,32 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/astutils"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 	"github.com/sqlc-dev/sqlc/internal/sql/validate"
 	"github.com/sqlc-dev/sqlc/internal/sqlcdebug"
 )
 
 var debugDumpAST = sqlcdebug.New("dumpast")
 
-func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query, error) {
+func (c *Compiler) parseQuery(stmt ast.Node, pp *preprocess.Result, o opts.Parser) (*Query, error) {
+	raw, ok := stmt.(*ast.RawStmt)
+	if !ok {
+		return nil, errors.New("node is not a statement")
+	}
+
+	// The preprocessor already replaced every sqlc.* construct with native SQL
+	// and recorded what it replaced.
+	pre := pp.Statement(raw.StmtLocation)
+	if pre.Err != nil {
+		return nil, pre.Err
+	}
+
+	// Engines number bind parameters in whatever order they convert the AST.
+	// Restore the numbering the preprocessor assigned in source order.
+	renumberParams(raw, pre.Numbers)
+
 	if c.coreCatalog != nil {
-		return c.parseQueryCore(stmt, src)
+		return c.parseQueryCore(raw, pp.Text, pre)
 	}
 
 	ctx := context.Background()
@@ -29,18 +46,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		debug.Dump(stmt)
 	}
 
-	// validate sqlc-specific syntax
-	if err := validate.SqlcFunctions(stmt); err != nil {
-		return nil, err
-	}
-
-	// rewrite queries to remove sqlc.* functions
-
-	raw, ok := stmt.(*ast.RawStmt)
-	if !ok {
-		return nil, errors.New("node is not a statement")
-	}
-	rawSQL, err := source.Pluck(src, raw.StmtLocation, raw.StmtLen)
+	rawSQL, err := source.Pluck(pp.Text, raw.StmtLocation, raw.StmtLen)
 	if err != nil {
 		return nil, err
 	}
@@ -128,7 +134,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 			Query:      expandedQuery,
 		}
 	} else if c.analyzer != nil {
-		inference, _ := c.inferQuery(raw, rawSQL)
+		inference, _ := c.inferQuery(raw, rawSQL, pre)
 		if inference == nil {
 			inference = &analysis{}
 		}
@@ -156,7 +162,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		// FOOTGUN: combineAnalysis mutates inference
 		anlys = combineAnalysis(inference, result)
 	} else {
-		anlys, err = c.analyzeQuery(raw, rawSQL)
+		anlys, err = c.analyzeQuery(raw, rawSQL, pre)
 		if err != nil {
 			return nil, err
 		}
@@ -188,6 +194,23 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 	}, nil
 }
 
+// renumberParams applies the parameter numbers the preprocessor assigned,
+// matching each node by the location of its placeholder in the query text.
+func renumberParams(raw *ast.RawStmt, numbers map[int]int) {
+	if len(numbers) == 0 {
+		return
+	}
+	astutils.Walk(astutils.VisitorFunc(func(node ast.Node) {
+		ref, ok := node.(*ast.ParamRef)
+		if !ok {
+			return
+		}
+		if n, ok := numbers[ref.Location]; ok {
+			ref.Number = n
+		}
+	}), raw)
+}
+
 func rangeVars(root ast.Node) []*ast.RangeVar {
 	var vars []*ast.RangeVar
 	find := astutils.VisitorFunc(func(node ast.Node) {

internal/compiler/parse_core.go

@@ -10,15 +10,11 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/named"
-	"github.com/sqlc-dev/sqlc/internal/sql/rewrite"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 	"github.com/sqlc-dev/sqlc/internal/sql/validate"
 )
 
-func (c *Compiler) parseQueryCore(stmt ast.Node, src string) (*Query, error) {
-	raw, ok := stmt.(*ast.RawStmt)
-	if !ok {
-		return nil, errors.New("node is not a statement")
-	}
+func (c *Compiler) parseQueryCore(raw *ast.RawStmt, src string, pre *preprocess.Statement) (*Query, error) {
 	rawSQL, err := source.Pluck(src, raw.StmtLocation, raw.StmtLen)
 	if err != nil {
 		return nil, err
@@ -48,21 +44,17 @@ func (c *Compiler) parseQueryCore(stmt ast.Node, src string) (*Query, error) {
 		return nil, err
 	}
 
-	numbers, dollar, err := validate.ParamRef(raw)
-	if err != nil {
-		return nil, err
-	}
-	rewritten, namedParams, edits := rewrite.NamedParameters(c.conf.Engine, raw, numbers, dollar)
-	expanded, err := source.Mutate(rawSQL, edits)
-	if err != nil {
-		return nil, err
+	if pre.ParamErr != nil {
+		return nil, pre.ParamErr
 	}
+	namedParams := pre.Params
+	expanded := rawSQL
 
 	var cols []*Column
 	var params []Parameter
-	switch rewritten.Stmt.(type) {
+	switch raw.Stmt.(type) {
 	case *ast.SelectStmt, *ast.InsertStmt, *ast.UpdateStmt, *ast.DeleteStmt:
-		res, err := coreanalyzer.Prepare(c.coreCatalog, rewritten)
+		res, err := coreanalyzer.Prepare(c.coreCatalog, raw)
 		if err != nil {
 			return nil, err
 		}
@@ -123,12 +115,17 @@ func coreParamColumn(p core.Parameter, params *named.ParamSet) *Column {
 		col.Table = &ast.TableName{Schema: p.Source.Schema, Name: p.Source.Table}
 		col.OriginalName = p.Source.Column
 	}
-	if col.Name == "" {
-		if name, ok := params.NameFor(p.Number); ok && name != "" {
-			col.Name = name
-		} else if p.Source != nil {
-			col.Name = p.Source.Column
-		}
+	if col.Name == "" && p.Source != nil {
+		col.Name = p.Source.Column
+	}
+	// Merge in what the user asked for: sqlc.narg() makes the parameter
+	// nullable and sqlc.slice() marks it as a slice, whichever way the
+	// analyzer typed it.
+	if param, isNamed := params.FetchMerge(p.Number, named.NewInferredParam(col.Name, p.NotNull)); isNamed {
+		col.Name = param.Name()
+		col.NotNull = param.NotNull()
+		col.IsSqlcSlice = param.IsSqlcSlice()
+		col.IsNamedParam = true
 	}
 	return col
 }

internal/compiler/query_catalog.go

@@ -5,16 +5,16 @@ import (
 
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/catalog"
-	"github.com/sqlc-dev/sqlc/internal/sql/rewrite"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 )
 
 type QueryCatalog struct {
 	catalog *catalog.Catalog
 	ctes    map[string]*Table
-	embeds  rewrite.EmbedSet
+	embeds  preprocess.EmbedSet
 }
 
-func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embeds rewrite.EmbedSet) (*QueryCatalog, error) {
+func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embeds preprocess.EmbedSet) (*QueryCatalog, error) {
 	var with *ast.WithClause
 	switch n := node.(type) {
 	case *ast.DeleteStmt:

internal/compiler/resolve.go

@@ -9,7 +9,7 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/sql/astutils"
 	"github.com/sqlc-dev/sqlc/internal/sql/catalog"
 	"github.com/sqlc-dev/sqlc/internal/sql/named"
-	"github.com/sqlc-dev/sqlc/internal/sql/rewrite"
+	"github.com/sqlc-dev/sqlc/internal/sql/preprocess"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlerr"
 )
 
@@ -21,7 +21,7 @@ func dataType(n *ast.TypeName) string {
 	}
 }
 
-func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar, args []paramRef, params *named.ParamSet, embeds rewrite.EmbedSet) ([]Parameter, error) {
+func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar, args []paramRef, params *named.ParamSet, embeds preprocess.EmbedSet) ([]Parameter, error) {
 	c := comp.catalog
 
 	aliasMap := map[string]*ast.TableName{}