Support :many key=group_id to return a map instead of a slice

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,18 @@ 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,
+				NoInference:    q.CmdParams.NoInference,
+			}
+		}
 		out = append(out, &plugin.Query{
 			Name:            q.Name,
 			Cmd:             q.Cmd,
+			CmdParams:       cmdParams,
 			Text:            q.SQL,
 			Comments:        q.Comments,
 			Columns:         columns,

internal/codegen/golang/query.go

@@ -200,6 +200,8 @@ type Query struct {
 	SourceName   string
 	Ret          QueryValue
 	Arg          QueryValue
+	ManyKey      string // Taken from CmdParams.ManyKey, the map-key for :many. If empty a slice should be returned.
+	ManyKeyType  string // The Go type of ManyKey
 	// Used for :copyfrom
 	Table *plugin.Identifier
 }
@@ -219,3 +221,23 @@ func (q Query) TableIdentifier() string {
 	}
 	return "[]string{" + strings.Join(escapedNames, ", ") + "}"
 }
+
+func (v Query) DefineRetTypeMultiple() string {
+	if v.ManyKey != "" {
+		return "map[" + v.ManyKeyType + "]" + v.Ret.DefineType()
+	}
+	return "[]" + v.Ret.DefineType()
+}
+
+func (v Query) HasManyKey() bool {
+	return v.ManyKey != ""
+}
+
+func (v Query) ManyKeyField() string {
+	for _, f := range v.Ret.Struct.Fields {
+		if f.DBName == v.ManyKey {
+			return v.Ret.Name + "." + f.Name
+		}
+	}
+	panic("couldn't find :many-key in struct fields")
+}

internal/codegen/golang/result.go

@@ -247,6 +247,9 @@ func buildQueries(req *plugin.CodeGenRequest, structs []Struct) ([]Query, error)
 		}
 
 		if len(query.Columns) == 1 && query.Columns[0].EmbedTable == nil {
+			if query.CmdParams.GetManyKey() != "" {
+				return nil, fmt.Errorf(":many key=%s query %s has only one column", query.CmdParams.GetManyKey(), query.Name)
+			}
 			c := query.Columns[0]
 			name := columnName(c, 0)
 			if c.IsFuncCall {
@@ -305,6 +308,18 @@ func buildQueries(req *plugin.CodeGenRequest, structs []Struct) ([]Query, error)
 				SQLDriver:   sqlpkg,
 				EmitPointer: req.Settings.Go.EmitResultStructPointers,
 			}
+			if query.CmdParams.GetManyKey() != "" {
+				gq.ManyKey = query.CmdParams.GetManyKey()
+				for _, c := range gs.Fields {
+					if c.DBName == gq.ManyKey {
+						gq.ManyKeyType = c.Type
+						break
+					}
+				}
+				if gq.ManyKeyType == "" {
+					return nil, fmt.Errorf("can not find key column %q in query %s", gq.ManyKey, gq.MethodName)
+				}
+			}
 		}
 
 		qs = append(qs, gq)

internal/codegen/golang/templates/stdlib/queryCode.tmpl

@@ -35,23 +35,27 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}
 {{if eq .Cmd ":many"}}
 {{range .Comments}}//{{.}}
 {{end -}}
-func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {
+func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}) ({{.DefineRetTypeMultiple}}, error) {
     {{- template "queryCodeStdExec" . }}
     if err != nil {
         return nil, err
     }
     defer rows.Close()
-    {{- if $.EmitEmptySlices}}
-    items := []{{.Ret.DefineType}}{}
+    {{- if (or $.EmitEmptySlices .HasManyKey)}}
+    items := {{.DefineRetTypeMultiple}}{}
     {{else}}
-    var items []{{.Ret.DefineType}}
+    var items {{.DefineRetTypeMultiple}}
     {{end -}}
     for rows.Next() {
         var {{.Ret.Name}} {{.Ret.Type}}
         if err := rows.Scan({{.Ret.Scan}}); err != nil {
             return nil, err
         }
+        {{- if .HasManyKey}}
+        items[{{.ManyKeyField}}] = {{.Ret.ReturnName}}
+        {{else}}
         items = append(items, {{.Ret.ReturnName}})
+        {{end -}}
     }
     if err := rows.Close(); err != nil {
         return nil, err

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
 }
 

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

internal/endtoend/testdata/codegen_json/gen/codegen.json

@@ -62583,6 +62583,7 @@
       "text": "SELECT id, name, bio FROM authors\nWHERE id = $1 LIMIT 1",
       "name": "GetAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62698,6 +62699,7 @@
       "text": "SELECT id, name, bio FROM authors\nORDER BY name",
       "name": "ListAuthors",
       "cmd": ":many",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62784,6 +62786,7 @@
       "text": "INSERT INTO authors (\n          name, bio\n) VALUES (\n  $1, $2\n)\nRETURNING id, name, bio",
       "name": "CreateAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62931,6 +62934,7 @@
       "text": "DELETE FROM authors\nWHERE id = $1",
       "name": "DeleteAuthor",
       "cmd": ":exec",
+      "cmd_params": null,
       "columns": [],
       "params": [
         {

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

@@ -1,6 +1,6 @@
 # package querytest
-query.sql:4:1: invalid query comment: -- name: ListFoos
-query.sql:7:1: invalid query comment: -- name: ListFoos :one :many
+query.sql:4:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
+query.sql:7:1: invalid query command parameter ":many"
 query.sql:10:1: invalid query type: :two
 query.sql:13:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause
 query.sql:16:1: query "UpdateFoo" specifies parameter ":one" without containing a RETURNING clause

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

@@ -1,6 +1,6 @@
 # package querytest
-query.sql:4:1: invalid query comment: -- name: ListFoos
-query.sql:7:1: invalid query comment: -- name: ListFoos :one :many
+query.sql:4:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
+query.sql:7:1: invalid query command parameter ":many"
 query.sql:10:1: invalid query type: :two
 query.sql:13:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause
 query.sql:16:1: query "UpdateFoo" specifies parameter ":one" without containing a RETURNING clause

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

@@ -1,6 +1,6 @@
 # package querytest
-query.sql:4:1: invalid query comment: -- name: ListFoos
-query.sql:7:1: invalid query comment: -- name: ListFoos :one :many
+query.sql:4:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
+query.sql:7:1: invalid query command parameter ":many"
 query.sql:10:1: invalid query type: :two
 query.sql:13:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause
 query.sql:16:1: query "UpdateFoo" specifies parameter ":one" without containing a RETURNING clause

internal/endtoend/testdata/manykey/stdlib/query.sql

@@ -0,0 +1,7 @@
+CREATE TABLE foo (
+	group_id INT NOT NULL,
+	score INT NOT NULL
+);
+
+-- name: SelectScoreSums :many key=group_id
+SELECT group_id, SUM(score) FROM foo GROUP BY group_id;

internal/endtoend/testdata/manykey/stdlib/sqlc.json

@@ -0,0 +1,11 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql"
+    }
+  ]
+}

internal/endtoend/testdata/process_plugin_disabled/gen/codegen.json

@@ -62583,6 +62583,7 @@
       "text": "SELECT id, name, bio FROM authors\nWHERE id = $1 LIMIT 1",
       "name": "GetAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62698,6 +62699,7 @@
       "text": "SELECT id, name, bio FROM authors\nORDER BY name",
       "name": "ListAuthors",
       "cmd": ":many",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62784,6 +62786,7 @@
       "text": "INSERT INTO authors (\n          name, bio\n) VALUES (\n  $1, $2\n)\nRETURNING id, name, bio",
       "name": "CreateAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62931,6 +62934,7 @@
       "text": "DELETE FROM authors\nWHERE id = $1",
       "name": "DeleteAuthor",
       "cmd": ":exec",
+      "cmd_params": null,
       "columns": [],
       "params": [
         {

internal/endtoend/testdata/process_plugin_sqlc_gen_json/gen/codegen.json

@@ -62583,6 +62583,7 @@
       "text": "SELECT id, name, bio FROM authors\nWHERE id = $1 LIMIT 1",
       "name": "GetAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62698,6 +62699,7 @@
       "text": "SELECT id, name, bio FROM authors\nORDER BY name",
       "name": "ListAuthors",
       "cmd": ":many",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62784,6 +62786,7 @@
       "text": "INSERT INTO authors (\n          name, bio\n) VALUES (\n  $1, $2\n)\nRETURNING id, name, bio",
       "name": "CreateAuthor",
       "cmd": ":one",
+      "cmd_params": null,
       "columns": [
         {
           "name": "id",
@@ -62931,6 +62934,7 @@
       "text": "DELETE FROM authors\nWHERE id = $1",
       "name": "DeleteAuthor",
       "cmd": ":exec",
+      "cmd_params": null,
       "columns": [],
       "params": [
         {

internal/metadata/meta.go

@@ -25,6 +25,12 @@ const (
 	CmdBatchOne   = ":batchone"
 )
 
+type CmdParams struct {
+	ManyKey        string
+	InsertMultiple bool
+	NoInference    bool
+}
+
 // A query name must be a valid Go identifier
 //
 // https://golang.org/ref/spec#Identifiers
@@ -44,7 +50,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 +82,56 @@ 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)
+		if len(part) < 4 {
+			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 p == "no-inference" {
+			if queryType != CmdExec && queryType != CmdExecResult && queryType != CmdExecRows && queryType != CmdExecLastId {
+				return ret, fmt.Errorf("query command parameter no-inference is invalid for query type %s", queryType)
+			}
+			ret.NoInference = 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
 }

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)
+			}
+		})
 	}
-
 }

protos/plugin/codegen.proto

@@ -179,13 +179,21 @@ message Query
   string text = 1 [json_name="text"];
   string name = 2 [json_name="name"];
   string cmd = 3 [json_name="cmd"];
+  CmdParams cmd_params = 9 [json_name="cmd_params"];
   repeated Column columns = 4 [json_name="columns"];
   repeated Parameter params = 5 [json_name="parameters"];
   repeated string comments = 6 [json_name="comments"];
   string filename = 7 [json_name="filename"];
   Identifier insert_into_table = 8 [json_name="insert_into_table"];
 }
 
+message CmdParams
+{
+  string many_key = 1 [json_name="many_key"];
+  bool insert_multiple = 2 [json_name="insert_multiple"];
+  bool no_inference = 3 [json_name="no_inference"];
+}
+
 message Parameter
 {
   int32 number = 1 [json_name="number"];
@@ -225,4 +233,4 @@ message VetQuery
   string name = 2 [json_name="name"];
   string cmd = 3 [json_name="cmd"];
   repeated VetParameter params = 4 [json_name="parameters"];
-}
+}