feat(compiler): Parse query parameter metadata from comments (#2850)

* feat(compiler): Parse query parameter metadata from comments

* update comment stripping, fix a test

* parse `@param` name correctly

* add a test for metadata param parsing

* use a word scanner for flag and param comment line parsing

* forgot that I partially cleaned comments already

* slightly more useful metadata test output

* drop unnecessary metadata test cases

* walk back changes, implement comment cleaning in source package

* cleanup

* simplify diff

* ensure at least one test case for parsing @whatever without a leading space

Author: andrewmbenton

internal/cmd/shim.go

@@ -209,13 +209,13 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
 			}
 		}
 		out = append(out, &plugin.Query{
-			Name:            q.Name,
-			Cmd:             q.Cmd,
+			Name:            q.Metadata.Name,
+			Cmd:             q.Metadata.Cmd,
 			Text:            q.SQL,
-			Comments:        q.Comments,
+			Comments:        q.Metadata.Comments,
 			Columns:         columns,
 			Params:          params,
-			Filename:        q.Filename,
+			Filename:        q.Metadata.Filename,
 			InsertIntoTable: iit,
 		})
 	}

internal/cmd/vet.go

@@ -545,7 +545,7 @@ func (c *checker) checkSQL(ctx context.Context, s config.SQL) error {
 	req := codeGenRequest(result, combo)
 	cfg := vetConfig(req)
 	for i, query := range req.Queries {
-		if result.Queries[i].Flags[QueryFlagSqlcVetDisable] {
+		if result.Queries[i].Metadata.Flags[QueryFlagSqlcVetDisable] {
 			if debug.Active {
 				log.Printf("Skipping vet rules for query: %s\n", query.Name)
 			}

internal/compiler/compile.go

@@ -8,10 +8,10 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/sqlc-dev/sqlc/internal/metadata"
 	"github.com/sqlc-dev/sqlc/internal/migrations"
 	"github.com/sqlc-dev/sqlc/internal/multierr"
 	"github.com/sqlc-dev/sqlc/internal/opts"
+	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlerr"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
@@ -20,7 +20,7 @@ import (
 // TODO: Rename this interface Engine
 type Parser interface {
 	Parse(io.Reader) ([]ast.Statement, error)
-	CommentSyntax() metadata.CommentSyntax
+	CommentSyntax() source.CommentSyntax
 	IsReservedKeyword(string) bool
 }
 
@@ -90,14 +90,15 @@ func (c *Compiler) parseQueries(o opts.Parser) (*Result, error) {
 				merr.Add(filename, src, loc, err)
 				continue
 			}
-			if query.Name != "" {
-				if _, exists := set[query.Name]; exists {
-					merr.Add(filename, src, stmt.Raw.Pos(), fmt.Errorf("duplicate query name: %s", query.Name))
+			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))
 					continue
 				}
-				set[query.Name] = struct{}{}
+				set[queryName] = struct{}{}
 			}
-			query.Filename = filepath.Base(filename)
+			query.Metadata.Filename = filepath.Base(filename)
 			if query != nil {
 				q = append(q, query)
 			}

internal/compiler/parse.go

@@ -43,14 +43,31 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		return nil, errors.New("missing semicolon at end of file")
 	}
 
-	name, cmd, err := metadata.ParseQueryNameAndType(strings.TrimSpace(rawSQL), c.parser.CommentSyntax())
+	name, cmd, err := metadata.ParseQueryNameAndType(rawSQL, metadata.CommentSyntax(c.parser.CommentSyntax()))
 	if err != nil {
 		return nil, err
 	}
+
 	if err := validate.Cmd(raw.Stmt, name, cmd); err != nil {
 		return nil, err
 	}
 
+	md := metadata.Metadata{
+		Name: name,
+		Cmd:  cmd,
+	}
+
+	// TODO eventually can use this for name and type/cmd parsing too
+	cleanedComments, err := source.CleanedComments(rawSQL, c.parser.CommentSyntax())
+	if err != nil {
+		return nil, err
+	}
+
+	md.Params, md.Flags, err = metadata.ParseParamsAndFlags(cleanedComments)
+	if err != nil {
+		return nil, err
+	}
+
 	var anlys *analysis
 	if c.analyzer != nil {
 		// TODO: Handle panics
@@ -90,17 +107,11 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		return nil, err
 	}
 
-	flags, err := metadata.ParseQueryFlags(comments)
-	if err != nil {
-		return nil, err
-	}
+	md.Comments = comments
 
 	return &Query{
 		RawStmt:         raw,
-		Cmd:             cmd,
-		Comments:        comments,
-		Name:            name,
-		Flags:           flags,
+		Metadata:        md,
 		Params:          anlys.Parameters,
 		Columns:         anlys.Columns,
 		SQL:             trimmed,

internal/compiler/query.go

@@ -1,6 +1,7 @@
 package compiler
 
 import (
+	"github.com/sqlc-dev/sqlc/internal/metadata"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 )
 
@@ -41,15 +42,9 @@ type Column struct {
 
 type Query struct {
 	SQL      string
-	Name     string
-	Cmd      string // TODO: Pick a better name. One of: one, many, exec, execrows, copyFrom
-	Flags    map[string]bool
+	Metadata metadata.Metadata
 	Columns  []*Column
 	Params   []Parameter
-	Comments []string
-
-	// XXX: Hack
-	Filename string
 
 	// Needed for CopyFrom
 	InsertIntoTable *ast.TableName

internal/endtoend/testdata/comment_syntax/mysql/go/query.sql.go

@@ -1,5 +1,5 @@
 # package querytest
-query.sql:1:1: invalid query comment: -- name: ListFoos
+query.sql:1:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
 query.sql:5:1: invalid query comment: -- name: ListFoos :one :many
 query.sql:8:1: invalid query type: :two
 query.sql:11:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause

internal/endtoend/testdata/invalid_queries_foo/pgx/v4/stderr.txt

@@ -1,5 +1,5 @@
 # package querytest
-query.sql:1:1: invalid query comment: -- name: ListFoos
+query.sql:1:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
 query.sql:5:1: invalid query comment: -- name: ListFoos :one :many
 query.sql:8:1: invalid query type: :two
 query.sql:11:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause

internal/endtoend/testdata/invalid_queries_foo/pgx/v5/stderr.txt

@@ -1,5 +1,5 @@
 # package querytest
-query.sql:1:1: invalid query comment: -- name: ListFoos
+query.sql:1:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
 query.sql:5:1: invalid query comment: -- name: ListFoos :one :many
 query.sql:8:1: invalid query type: :two
 query.sql:11:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause

internal/endtoend/testdata/invalid_queries_foo/stdlib/stderr.txt

@@ -10,7 +10,7 @@ import (
 	"github.com/pingcap/tidb/parser"
 	_ "github.com/pingcap/tidb/parser/test_driver"
 
-	"github.com/sqlc-dev/sqlc/internal/metadata"
+	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlerr"
 )
@@ -86,8 +86,8 @@ func (p *Parser) Parse(r io.Reader) ([]ast.Statement, error) {
 }
 
 // https://dev.mysql.com/doc/refman/8.0/en/comments.html
-func (p *Parser) CommentSyntax() metadata.CommentSyntax {
-	return metadata.CommentSyntax{
+func (p *Parser) CommentSyntax() source.CommentSyntax {
+	return source.CommentSyntax{
 		Dash:      true,
 		SlashStar: true,
 		Hash:      true,

internal/engine/dolphin/parse.go

@@ -12,7 +12,7 @@ import (
 	nodes "github.com/pganalyze/pg_query_go/v4"
 	"github.com/pganalyze/pg_query_go/v4/parser"
 
-	"github.com/sqlc-dev/sqlc/internal/metadata"
+	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlerr"
 )
@@ -199,8 +199,8 @@ func normalizeErr(err error) error {
 }
 
 // https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-COMMENTS
-func (p *Parser) CommentSyntax() metadata.CommentSyntax {
-	return metadata.CommentSyntax{
+func (p *Parser) CommentSyntax() source.CommentSyntax {
+	return source.CommentSyntax{
 		Dash:      true,
 		SlashStar: true,
 	}

internal/engine/postgresql/parse.go

@@ -8,7 +8,7 @@ import (
 	"github.com/antlr/antlr4/runtime/Go/antlr/v4"
 
 	"github.com/sqlc-dev/sqlc/internal/engine/sqlite/parser"
-	"github.com/sqlc-dev/sqlc/internal/metadata"
+	"github.com/sqlc-dev/sqlc/internal/source"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 )
 
@@ -86,8 +86,8 @@ func (p *Parser) Parse(r io.Reader) ([]ast.Statement, error) {
 	return stmts, nil
 }
 
-func (p *Parser) CommentSyntax() metadata.CommentSyntax {
-	return metadata.CommentSyntax{
+func (p *Parser) CommentSyntax() source.CommentSyntax {
+	return source.CommentSyntax{
 		Dash:      true,
 		Hash:      false,
 		SlashStar: true,

internal/engine/sqlite/parse.go

@@ -1,15 +1,24 @@
 package metadata
 
 import (
+	"bufio"
 	"fmt"
 	"strings"
 	"unicode"
+
+	"github.com/sqlc-dev/sqlc/internal/source"
 )
 
-type CommentSyntax struct {
-	Dash      bool
-	Hash      bool
-	SlashStar bool
+type CommentSyntax source.CommentSyntax
+
+type Metadata struct {
+	Name     string
+	Cmd      string
+	Comments []string
+	Params   map[string]string
+	Flags    map[string]bool
+
+	Filename string
 }
 
 const (
@@ -83,7 +92,7 @@ func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string
 		if prefix == "/*" {
 			part = part[:len(part)-1] // removes the trailing "*/" element
 		}
-		if len(part) == 2 {
+		if len(part) == 3 {
 			return "", "", fmt.Errorf("missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: %s", line)
 		}
 		if len(part) != 4 {
@@ -104,19 +113,39 @@ func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string
 	return "", "", nil
 }
 
-func ParseQueryFlags(comments []string) (map[string]bool, error) {
+func ParseParamsAndFlags(comments []string) (map[string]string, map[string]bool, error) {
+	params := make(map[string]string)
 	flags := make(map[string]bool)
+
 	for _, line := range comments {
-		cleanLine := strings.TrimPrefix(line, "--")
-		cleanLine = strings.TrimPrefix(cleanLine, "/*")
-		cleanLine = strings.TrimPrefix(cleanLine, "#")
-		cleanLine = strings.TrimSuffix(cleanLine, "*/")
-		cleanLine = strings.TrimSpace(cleanLine)
-		if strings.HasPrefix(cleanLine, "@") {
-			flagName := strings.SplitN(cleanLine, " ", 2)[0]
-			flags[flagName] = true
+		s := bufio.NewScanner(strings.NewReader(line))
+		s.Split(bufio.ScanWords)
+
+		s.Scan()
+		token := s.Text()
+
+		if !strings.HasPrefix(token, "@") {
 			continue
 		}
+
+		switch token {
+		case "@param":
+			s.Scan()
+			name := s.Text()
+			var rest []string
+			for s.Scan() {
+				paramToken := s.Text()
+				rest = append(rest, paramToken)
+			}
+			params[name] = strings.Join(rest, " ")
+		default:
+			flags[token] = true
+		}
+
+		if s.Err() != nil {
+			return params, flags, s.Err()
+		}
 	}
-	return flags, nil
+
+	return params, flags, nil
 }