From e38ddf16d111d55929c31719fbe7da1701fd1b43 Mon Sep 17 00:00:00 2001
From: Andrew Benton <andrew@sqlc.dev>
Date: Fri, 14 Jul 2023 17:10:05 -0400
Subject: [PATCH] Parse comment lines starting with @symbol as boolean flags
 associated with a query

---
 internal/compiler/parse.go     |  9 ++++++++-
 internal/compiler/query.go     |  1 +
 internal/metadata/meta.go      | 18 +++++++++++++++++-
 internal/metadata/meta_test.go | 26 ++++++++++++++++++++++----
 4 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/internal/compiler/parse.go b/internal/compiler/parse.go
index b108acc492..932df32d32 100644
--- a/internal/compiler/parse.go
+++ b/internal/compiler/parse.go
@@ -70,7 +70,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 	if err := validate.In(c.catalog, raw); err != nil {
 		return nil, err
 	}
-	name, cmd, err := metadata.Parse(strings.TrimSpace(rawSQL), c.parser.CommentSyntax())
+	name, cmd, err := metadata.ParseQueryNameAndType(strings.TrimSpace(rawSQL), c.parser.CommentSyntax())
 	if err != nil {
 		return nil, err
 	}
@@ -125,11 +125,18 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 	if err != nil {
 		return nil, err
 	}
+
+	flags, err := metadata.ParseQueryFlags(comments)
+	if err != nil {
+		return nil, err
+	}
+
 	return &Query{
 		RawStmt:         raw,
 		Cmd:             cmd,
 		Comments:        comments,
 		Name:            name,
+		Flags:           flags,
 		Params:          params,
 		Columns:         cols,
 		SQL:             trimmed,
diff --git a/internal/compiler/query.go b/internal/compiler/query.go
index e77f555dbd..9a2df0351c 100644
--- a/internal/compiler/query.go
+++ b/internal/compiler/query.go
@@ -42,6 +42,7 @@ 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
 	Columns  []*Column
 	Params   []Parameter
 	Comments []string
diff --git a/internal/metadata/meta.go b/internal/metadata/meta.go
index 59a1add763..3c34eff58a 100644
--- a/internal/metadata/meta.go
+++ b/internal/metadata/meta.go
@@ -44,7 +44,7 @@ func validateQueryName(name string) error {
 	return nil
 }
 
-func Parse(t string, commentStyle CommentSyntax) (string, string, error) {
+func ParseQueryNameAndType(t string, commentStyle CommentSyntax) (string, string, error) {
 	for _, line := range strings.Split(t, "\n") {
 		var prefix string
 		if strings.HasPrefix(line, "--") {
@@ -103,3 +103,19 @@ func Parse(t string, commentStyle CommentSyntax) (string, string, error) {
 	}
 	return "", "", nil
 }
+
+func ParseQueryFlags(comments []string) (map[string]bool, error) {
+	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
+		}
+	}
+	return flags, nil
+}
diff --git a/internal/metadata/meta_test.go b/internal/metadata/meta_test.go
index e2f7905cba..cbfcb6fba6 100644
--- a/internal/metadata/meta_test.go
+++ b/internal/metadata/meta_test.go
@@ -2,7 +2,7 @@ package metadata
 
 import "testing"
 
-func TestParseMetadata(t *testing.T) {
+func TestParseQueryNameAndType(t *testing.T) {
 
 	for _, query := range []string{
 		`-- name: CreateFoo, :one`,
@@ -17,7 +17,7 @@ func TestParseMetadata(t *testing.T) {
 		"-- name:CreateFoo",
 		`--name:CreateFoo :two`,
 	} {
-		if _, _, err := Parse(query, CommentSyntax{Dash: true}); err == nil {
+		if _, _, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true}); err == nil {
 			t.Errorf("expected invalid metadata: %q", query)
 		}
 	}
@@ -27,13 +27,13 @@ func TestParseMetadata(t *testing.T) {
 		`-- name comment`,
 		`--name comment`,
 	} {
-		if _, _, err := Parse(query, CommentSyntax{Dash: true}); err != nil {
+		if _, _, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true}); err != nil {
 			t.Errorf("expected valid comment: %q", query)
 		}
 	}
 
 	query := `-- name: CreateFoo :one`
-	queryName, queryType, err := Parse(query, CommentSyntax{Dash: true})
+	queryName, queryType, err := ParseQueryNameAndType(query, CommentSyntax{Dash: true})
 	if err != nil {
 		t.Errorf("expected valid metadata: %q", query)
 	}
@@ -45,3 +45,21 @@ func TestParseMetadata(t *testing.T) {
 	}
 
 }
+
+func TestParseQueryFlags(t *testing.T) {
+	for _, comments := range [][]string{
+		{
+			"-- name: CreateFoo :one",
+			"-- @flag-foo",
+		},
+	} {
+		flags, err := ParseQueryFlags(comments)
+		if err != nil {
+			t.Errorf("expected query flags to parse, got error: %s", err)
+		}
+
+		if !flags["@flag-foo"] {
+			t.Errorf("expected flag not found")
+		}
+	}
+}
\ No newline at end of file