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": [
         {