feat(compiler): Parse query parameter metadata from comments

sqlc-dev · andrewmbenton · Oct 16, 2023 · Oct 12, 2023 · Oct 13, 2023 · Oct 13, 2023
commit 4db5c66374eee141029c7e7af70b7a70aa05de3c
diff --git a/internal/cmd/shim.go b/internal/cmd/shim.go
@@ -181,13 +181,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,
 		})
 	}

diff --git a/internal/cmd/vet.go b/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)
 			}

diff --git a/internal/compiler/compile.go b/internal/compiler/compile.go
@@ -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)
 			}

diff --git a/internal/compiler/parse.go b/internal/compiler/parse.go
@@ -43,11 +43,12 @@ 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())
+	md, err := metadata.ParseQueryMetadata(rawSQL, c.parser.CommentSyntax())
 	if err != nil {
 		return nil, err
 	}
-	if err := validate.Cmd(raw.Stmt, name, cmd); err != nil {
+
+	if err := validate.Cmd(raw.Stmt, md.Name, md.Cmd); err != nil {
 		return nil, err
 	}
 
@@ -85,22 +86,14 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		}
 	}
 
-	trimmed, comments, err := source.StripComments(expanded)
-	if err != nil {
-		return nil, err
-	}
-
-	flags, err := metadata.ParseQueryFlags(comments)
+	trimmed, _, err := source.StripComments(expanded)
 	if err != nil {
 		return nil, err
 	}
 
 	return &Query{
 		RawStmt:         raw,
-		Cmd:             cmd,
-		Comments:        comments,
-		Name:            name,
-		Flags:           flags,
+		Metadata:        md,
 		Params:          anlys.Parameters,
 		Columns:         anlys.Columns,
 		SQL:             trimmed,

diff --git a/internal/compiler/query.go b/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

diff --git a/internal/endtoend/testdata/invalid_queries_foo/pgx/v4/stderr.txt b/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

diff --git a/internal/endtoend/testdata/invalid_queries_foo/pgx/v5/stderr.txt b/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

diff --git a/internal/endtoend/testdata/invalid_queries_foo/stdlib/stderr.txt b/internal/endtoend/testdata/invalid_queries_foo/stdlib/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

diff --git a/internal/metadata/meta.go b/internal/metadata/meta.go
@@ -1,11 +1,22 @@
 package metadata
 
 import (
+	"bufio"
 	"fmt"
 	"strings"
 	"unicode"
 )
 
+type Metadata struct {
+	Name     string
+	Cmd      string
+	Comments []string
+	Params   map[string]string
+	Flags    map[string]bool
+
+	Filename string
+}
+
 type CommentSyntax struct {
 	Dash      bool
 	Hash      bool
@@ -44,8 +55,12 @@ func validateQueryName(name string) error {
 	return nil
 }
 
-func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string, error) {
-	for _, line := range strings.Split(t, "\n") {
+func ParseQueryMetadata(rawSql string, commentStyle CommentSyntax) (Metadata, error) {
+	md := Metadata{}
+	s := bufio.NewScanner(strings.NewReader(strings.TrimSpace(rawSql)))
+	var lines, comments []string
+	for s.Scan() {
+		line := s.Text()
 		var prefix string
 		if strings.HasPrefix(line, "--") {
 			if !commentStyle.Dash {
@@ -66,45 +81,56 @@ func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string
 			prefix = "#"
 		}
 		if prefix == "" {
+			lines = append(lines, line)
 			continue
 		}
+
 		rest := line[len(prefix):]
+		rest = strings.TrimSuffix(rest, "*/")
+		comments = append(comments, rest)
+
 		if !strings.HasPrefix(strings.TrimSpace(rest), "name") {
 			continue
 		}
 		if !strings.Contains(rest, ":") {
 			continue
 		}
 		if !strings.HasPrefix(rest, " name: ") {
-			return "", "", fmt.Errorf("invalid metadata: %s", line)
+			return md, fmt.Errorf("invalid metadata: %s", line)
 		}
 
-		part := strings.Split(strings.TrimSpace(line), " ")
-		if prefix == "/*" {
-			part = part[:len(part)-1] // removes the trailing "*/" element
-		}
-		if len(part) == 2 {
-			return "", "", fmt.Errorf("missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: %s", line)
+		comments = comments[:len(comments)-1] // Remove tha name line from returned comments
+
+		parts := strings.Split(strings.TrimSpace(rest), " ")
+
+		if len(parts) == 2 {
+			return md, fmt.Errorf("missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: %s", line)
 		}
-		if len(part) != 4 {
-			return "", "", fmt.Errorf("invalid query comment: %s", line)
+		if len(parts) > 3 {
+			return md, fmt.Errorf("invalid query comment: %s", line)
 		}
-		queryName := part[2]
-		queryType := strings.TrimSpace(part[3])
+		queryName := parts[1]
+		queryType := parts[2]
 		switch queryType {
 		case CmdOne, CmdMany, CmdExec, CmdExecResult, CmdExecRows, CmdExecLastId, CmdCopyFrom, CmdBatchExec, CmdBatchMany, CmdBatchOne:
 		default:
-			return "", "", fmt.Errorf("invalid query type: %s", queryType)
+			return md, fmt.Errorf("invalid query type: %s", queryType)
 		}
 		if err := validateQueryName(queryName); err != nil {
-			return "", "", err
+			return md, err
 		}
-		return queryName, queryType, nil
+		md.Name = queryName
+		md.Cmd = queryType
 	}
-	return "", "", nil
+
+	md.Comments = comments
+	md.Params, md.Flags = parseParamsAndFlags(md.Comments)
+
+	return md, s.Err()
 }
 
-func ParseQueryFlags(comments []string) (map[string]bool, error) {
+func parseParamsAndFlags(comments []string) (map[string]string, map[string]bool) {
+	params := make(map[string]string)
 	flags := make(map[string]bool)
 	for _, line := range comments {
 		cleanLine := strings.TrimPrefix(line, "--")
@@ -113,10 +139,16 @@ func ParseQueryFlags(comments []string) (map[string]bool, error) {
 		cleanLine = strings.TrimSuffix(cleanLine, "*/")
 		cleanLine = strings.TrimSpace(cleanLine)
 		if strings.HasPrefix(cleanLine, "@") {
-			flagName := strings.SplitN(cleanLine, " ", 2)[0]
-			flags[flagName] = true
+			parts := strings.SplitN(cleanLine, " ", 2)
+			name := parts[0]
+			switch name {
+			case "param":
+				params[name] = parts[1]
+			default:
+				flags[name] = true
+			}
 			continue
 		}
 	}
-	return flags, nil
+	return params, flags
 }
diff --git a/internal/metadata/meta_test.go b/internal/metadata/meta_test.go
@@ -16,8 +16,11 @@ func TestParseQueryNameAndType(t *testing.T) {
 		`--name: CreateFoo :two`,
 		"-- name:CreateFoo",
 		`--name:CreateFoo :two`,
+		`--  name: CreateFoo :two`,
+		`-- name:  CreateFoo :two`,
+		`-- name: CreateFoo  :two`,
 	} {
-		if _, _, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true}); err == nil {
+		if _, err := ParseQueryMetadata(query, CommentSyntax{Dash: true}); err == nil {
 			t.Errorf("expected invalid metadata: %q", query)
 		}
 	}
@@ -27,21 +30,26 @@ func TestParseQueryNameAndType(t *testing.T) {
 		`-- name comment`,
 		`--name comment`,
 	} {
-		if _, _, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true}); err != nil {
+		if _, err := ParseQueryMetadata(query, CommentSyntax{Dash: true}); err != nil {
 			t.Errorf("expected valid comment: %q", query)
 		}
 	}
 
-	query := `-- name: CreateFoo :one`
-	queryName, queryType, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true})
-	if err != nil {
-		t.Errorf("expected valid metadata: %q", query)
-	}
-	if queryName != "CreateFoo" {
-		t.Errorf("incorrect queryName parsed: %q", query)
-	}
-	if queryType != CmdOne {
-		t.Errorf("incorrect queryType parsed: %q", query)
+	for query, cs := range map[string]CommentSyntax{
+		`-- name: CreateFoo :one`:    {Dash: true},
+		`# name: CreateFoo :one`:     {Hash: true},
+		`/* name: CreateFoo :one */`: {SlashStar: true},
+	} {
+		queryMetadata, err := ParseQueryMetadata(query, cs)
+		if err != nil {
+			t.Errorf("expected valid metadata: %q", query)
+		}
+		if queryMetadata.Name != "CreateFoo" {
+			t.Errorf("incorrect queryName parsed: %q", query)
+		}
+		if queryMetadata.Cmd != CmdOne {
+			t.Errorf("incorrect queryType parsed: %q", query)
+		}
 	}
 
 }
@@ -52,14 +60,24 @@ func TestParseQueryFlags(t *testing.T) {
 			"-- name: CreateFoo :one",
 			"-- @flag-foo",
 		},
+		{
+			"-- name: CreateFoo :one",
+			"-- @flag-foo @flag-bar",
+		},
+		{
+			"-- name: GetFoos :many",
+			"-- @param @flag-bar UUID",
+			"-- @flag-foo",
+		},
 	} {
-		flags, err := ParseQueryFlags(comments)
-		if err != nil {
-			t.Errorf("expected query flags to parse, got error: %s", err)
-		}
+		_, flags := parseParamsAndFlags(comments)
 
 		if !flags["@flag-foo"] {
 			t.Errorf("expected flag not found")
 		}
+
+		if flags["@flag-bar"] {
+			t.Errorf("unexpected flag found")
+		}
 	}
-}
+}
diff --git a/internal/source/code.go b/internal/source/code.go
@@ -90,6 +90,7 @@ func Mutate(raw string, a []Edit) (string, error) {
 	return s, nil
 }
 
+// TODO remove and roll into query parsing
 func StripComments(sql string) (string, []string, error) {
 	s := bufio.NewScanner(strings.NewReader(strings.TrimSpace(sql)))
 	var lines, comments []string