Allow for extra parameters to the queryType

Like `-- name: InsertMulti :exec multiple` and `-- name: Select :many key=group_id` issue sqlc-dev#2353
Jille · Jun 28, 2023 · 95f5175 · 95f5175
1 parent 3e0fca0
commit 95f5175
Show file tree

Hide file tree

Showing 8 changed files with 628 additions and 142 deletions.
diff --git a/internal/cmd/shim.go b/internal/cmd/shim.go
@@ -7,6 +7,7 @@ import (
 	"github.com/kyleconroy/sqlc/internal/config"
 	"github.com/kyleconroy/sqlc/internal/config/convert"
 	"github.com/kyleconroy/sqlc/internal/info"
+	"github.com/kyleconroy/sqlc/internal/metadata"
 	"github.com/kyleconroy/sqlc/internal/plugin"
 	"github.com/kyleconroy/sqlc/internal/sql/catalog"
 )
@@ -226,9 +227,17 @@ func pluginQueries(r *compiler.Result) []*plugin.Query {
 				Name:    q.InsertIntoTable.Name,
 			}
 		}
+		var cmdParams *plugin.CmdParams
+		if q.CmdParams != (metadata.CmdParams{}) {
+			cmdParams = &plugin.CmdParams{
+				ManyKey:        q.CmdParams.ManyKey,
+				InsertMultiple: q.CmdParams.InsertMultiple,
+			}
+		}
 		out = append(out, &plugin.Query{
 			Name:            q.Name,
 			Cmd:             q.Cmd,
+			CmdParams:       cmdParams,
 			Text:            q.SQL,
 			Comments:        q.Comments,
 			Columns:         columns,

diff --git 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, cmdParams, err := metadata.Parse(strings.TrimSpace(rawSQL), c.parser.CommentSyntax())
 	if err != nil {
 		return nil, err
 	}
@@ -133,6 +133,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 		Columns:         cols,
 		SQL:             trimmed,
 		InsertIntoTable: table,
+		CmdParams:       cmdParams,
 	}, nil
 }
 

diff --git a/internal/compiler/query.go b/internal/compiler/query.go
@@ -1,6 +1,7 @@
 package compiler
 
 import (
+	"github.com/kyleconroy/sqlc/internal/metadata"
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 )
 
@@ -39,12 +40,13 @@ type Column struct {
 }
 
 type Query struct {
-	SQL      string
-	Name     string
-	Cmd      string // TODO: Pick a better name. One of: one, many, exec, execrows, copyFrom
-	Columns  []*Column
-	Params   []Parameter
-	Comments []string
+	SQL       string
+	Name      string
+	Cmd       string // TODO: Pick a better name. One of: one, many, exec, execrows, copyFrom
+	CmdParams metadata.CmdParams
+	Columns   []*Column
+	Params    []Parameter
+	Comments  []string
 
 	// XXX: Hack
 	Filename string

diff --git a/internal/metadata/meta.go b/internal/metadata/meta.go
@@ -25,6 +25,11 @@ const (
 	CmdBatchOne   = ":batchone"
 )
 
+type CmdParams struct {
+	ManyKey        string
+	InsertMultiple bool
+}
+
 // A query name must be a valid Go identifier
 //
 // https://golang.org/ref/spec#Identifiers
@@ -44,7 +49,7 @@ func validateQueryName(name string) error {
 	return nil
 }
 
-func Parse(t string, commentStyle CommentSyntax) (string, string, error) {
+func Parse(t string, commentStyle CommentSyntax) (string, string, CmdParams, error) {
 	for _, line := range strings.Split(t, "\n") {
 		var prefix string
 		if strings.HasPrefix(line, "--") {
@@ -76,30 +81,51 @@ func Parse(t string, commentStyle CommentSyntax) (string, string, error) {
 			continue
 		}
 		if !strings.HasPrefix(rest, " name: ") {
-			return "", "", fmt.Errorf("invalid metadata: %s", line)
+			return "", "", CmdParams{}, 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)
-		}
-		if len(part) != 4 {
-			return "", "", fmt.Errorf("invalid query comment: %s", line)
+			return "", "", CmdParams{}, fmt.Errorf("missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: %s", line)
 		}
 		queryName := part[2]
 		queryType := strings.TrimSpace(part[3])
 		switch queryType {
 		case CmdOne, CmdMany, CmdExec, CmdExecResult, CmdExecRows, CmdExecLastId, CmdCopyFrom, CmdBatchExec, CmdBatchMany, CmdBatchOne:
 		default:
-			return "", "", fmt.Errorf("invalid query type: %s", queryType)
+			return "", "", CmdParams{}, fmt.Errorf("invalid query type: %s", queryType)
 		}
 		if err := validateQueryName(queryName); err != nil {
-			return "", "", err
+			return "", "", CmdParams{}, err
+		}
+		cmdParams, err := parseCmdParams(part[4:], queryType)
+		if err != nil {
+			return "", "", CmdParams{}, err
+		}
+		return queryName, queryType, cmdParams, nil
+	}
+	return "", "", CmdParams{}, nil
+}
+
+func parseCmdParams(part []string, queryType string) (CmdParams, error) {
+	var ret CmdParams
+	for _, p := range part {
+		if p == "multiple" {
+			if queryType != CmdExec && queryType != CmdExecResult && queryType != CmdExecRows && queryType != CmdExecLastId {
+				return ret, fmt.Errorf("query command parameter multiple is invalid for query type %s", queryType)
+			}
+			ret.InsertMultiple = true
+		} else if strings.HasPrefix(p, "key=") {
+			if queryType != CmdMany {
+				return ret, fmt.Errorf("query command parameter %s is invalid for query type %s", p, queryType)
+			}
+			ret.ManyKey = strings.TrimPrefix(p, "key=")
+		} else {
+			return ret, fmt.Errorf("invalid query command parameter %q", p)
 		}
-		return queryName, queryType, nil
 	}
-	return "", "", nil
+	return ret, nil
 }
diff --git 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 TestNonMetadata(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 := Parse(query, CommentSyntax{Dash: true}); err == nil {
 			t.Errorf("expected invalid metadata: %q", query)
 		}
 	}
@@ -27,21 +27,52 @@ func TestParseMetadata(t *testing.T) {
 		`-- name comment`,
 		`--name comment`,
 	} {
-		if _, _, err := Parse(query, CommentSyntax{Dash: true}); err != nil {
+		if _, _, _, err := Parse(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})
-	if err != nil {
-		t.Errorf("expected valid metadata: %q", query)
-	}
-	if queryName != "CreateFoo" {
-		t.Errorf("incorrect queryName parsed: %q", query)
+func TestParse(t *testing.T) {
+	tests := []struct {
+		query         string
+		wantName      string
+		wantType      string
+		wantCmdParams CmdParams
+	}{
+		{
+			query:    "-- name: CreateFoo :one",
+			wantName: "CreateFoo",
+			wantType: CmdOne,
+		},
+		{
+			query:         "-- name: InsertMulti :exec multiple",
+			wantName:      "InsertMulti",
+			wantType:      CmdExec,
+			wantCmdParams: CmdParams{InsertMultiple: true},
+		},
+		{
+			query:         "-- name: SelectKey :many key=group_id",
+			wantName:      "SelectKey",
+			wantType:      CmdMany,
+			wantCmdParams: CmdParams{ManyKey: "group_id"},
+		},
 	}
-	if queryType != CmdOne {
-		t.Errorf("incorrect queryType parsed: %q", query)
+	for _, tc := range tests {
+		t.Run(tc.query, func(t *testing.T) {
+			name, queryType, cmdParams, err := Parse(tc.query, CommentSyntax{Dash: true})
+			if err != nil {
+				t.Fatalf("Parse failed: %v", err)
+			}
+			if name != tc.wantName {
+				t.Errorf("unexpected name: got %q; want %q", name, tc.wantName)
+			}
+			if queryType != tc.wantType {
+				t.Errorf("unexpected queryType: got %q; want %q", queryType, tc.wantType)
+			}
+			if cmdParams != tc.wantCmdParams {
+				t.Errorf("unexpected cmdParams: got %#v; want %#v", cmdParams, tc.wantCmdParams)
+			}
+		})
 	}
-
 }