Skip to content

Commit e6bf21c

Browse files
jaglyserxjaglyserx
committed
feat(first): add new command for optional return 'first'
1 parent c048334 commit e6bf21c

File tree

14 files changed

+109
-9
lines changed

14 files changed

+109
-9
lines changed

‎internal/codegen/golang/gen.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,12 @@ func (t *tmplCtx) codegenQueryMethod(q Query) string {
7474
}
7575
return db + ".QueryRowContext"
7676

77+
case ":first":
78+
if t.EmitPreparedQueries {
79+
return "q.queryRow"
80+
}
81+
return db + ".QueryRowContext"
82+
7783
case ":many":
7884
if t.EmitPreparedQueries {
7985
return "q.query"
@@ -91,6 +97,8 @@ func (t *tmplCtx) codegenQueryMethod(q Query) string {
9197
func (t *tmplCtx) codegenQueryRetval(q Query) (string, error) {
9298
switch q.Cmd {
9399
case ":one":
100+
fallthrough
101+
case ":first":
94102
return "row :=", nil
95103
case ":many":
96104
return "rows, err :=", nil

‎internal/codegen/golang/imports.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,25 @@ func (i *importer) queryImports(filename string) fileImports {
342342
}
343343
return false
344344
})
345+
usesFirst := false
346+
for _, q := range gq {
347+
if q.Cmd == metadata.CmdFirst {
348+
usesFirst = true
349+
break
350+
}
351+
}
352+
if usesFirst {
353+
std["errors"] = struct{}{}
354+
sqlpkg := parseDriver(i.Options.SqlPackage)
355+
switch sqlpkg {
356+
case opts.SQLDriverPGXV4:
357+
pkg[ImportSpec{Path: "github.com/jackc/pgx/v4"}] = struct{}{}
358+
case opts.SQLDriverPGXV5:
359+
pkg[ImportSpec{Path: "github.com/jackc/pgx/v5"}] = struct{}{}
360+
default:
361+
std["database/sql"] = struct{}{}
362+
}
363+
}
345364

346365
sliceScan := func() bool {
347366
for _, q := range gq {

‎internal/codegen/golang/query.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
type QueryValue struct {
1313
Emit bool
1414
EmitPointer bool
15+
ForcePointer bool
1516
Name string
1617
DBName string // The name of the field in the database. Only set if Struct==nil.
1718
Struct *Struct
@@ -32,7 +33,7 @@ func (v QueryValue) IsStruct() bool {
3233
}
3334

3435
func (v QueryValue) IsPointer() bool {
35-
return v.EmitPointer && v.Struct != nil
36+
return v.ForcePointer || (v.EmitPointer && v.Struct != nil)
3637
}
3738

3839
func (v QueryValue) isEmpty() bool {
@@ -270,7 +271,7 @@ type Query struct {
270271
}
271272

272273
func (q Query) hasRetType() bool {
273-
scanned := q.Cmd == metadata.CmdOne || q.Cmd == metadata.CmdMany ||
274+
scanned := q.Cmd == metadata.CmdOne || q.Cmd == metadata.CmdFirst || q.Cmd == metadata.CmdMany ||
274275
q.Cmd == metadata.CmdBatchMany || q.Cmd == metadata.CmdBatchOne
275276
return scanned && !q.Ret.isEmpty()
276277
}

‎internal/codegen/golang/result.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -322,6 +322,9 @@ func buildQueries(req *plugin.GenerateRequest, options *opts.Options, structs []
322322
EmitPointer: options.EmitResultStructPointers,
323323
}
324324
}
325+
if query.Cmd == metadata.CmdFirst {
326+
gq.Ret.ForcePointer = true
327+
}
325328

326329
qs = append(qs, gq)
327330
}
@@ -334,6 +337,7 @@ var cmdReturnsData = map[string]struct{}{
334337
metadata.CmdBatchOne: {},
335338
metadata.CmdMany: {},
336339
metadata.CmdOne: {},
340+
metadata.CmdFirst: {},
337341
}
338342

339343
func putOutColumns(query *plugin.Query) bool {

‎internal/codegen/golang/result_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ func TestPutOutColumns_ForZeroColumns(t *testing.T) {
3636
cmd: metadata.CmdOne,
3737
want: true,
3838
},
39+
{
40+
cmd: metadata.CmdFirst,
41+
want: true,
42+
},
3943
{
4044
cmd: metadata.CmdCopyFrom,
4145
want: false,

‎internal/codegen/golang/templates/pgx/interfaceCode.tmpl‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,15 @@
1111
{{end -}}
1212
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
1313
{{- end}}
14+
{{- if and (eq .Cmd ":first") ($dbtxParam) }}
15+
{{range .Comments}}//{{.}}
16+
{{end -}}
17+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
18+
{{- else if eq .Cmd ":first" }}
19+
{{range .Comments}}//{{.}}
20+
{{end -}}
21+
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
22+
{{- end}}
1423
{{- if and (eq .Cmd ":many") ($dbtxParam) }}
1524
{{range .Comments}}//{{.}}
1625
{{end -}}

‎internal/codegen/golang/templates/pgx/queryCode.tmpl‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,31 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.De
4444
{{- end}}
4545
return {{.Ret.ReturnName}}, err
4646
}
47+
{{else if eq .Cmd ":first"}}
48+
{{range .Comments}}//{{.}}
49+
{{end -}}
50+
{{- if $.EmitMethodsWithDBArgument -}}
51+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
52+
row := db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
53+
{{- else -}}
54+
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
55+
row := q.db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
56+
{{- end}}
57+
{{- if or (ne .Arg.Pair .Ret.Pair) (ne .Arg.DefineType .Ret.DefineType) }}
58+
var {{.Ret.Name}} {{.Ret.Type}}
59+
{{- end}}
60+
err := row.Scan({{.Ret.Scan}})
61+
if err != nil {
62+
if errors.Is(err, pgx.ErrNoRows) {
63+
return nil, nil
64+
}
65+
{{- if $.WrapErrors}}
66+
err = fmt.Errorf("query {{.MethodName}}: %w", err)
67+
{{- end}}
68+
return nil, err
69+
}
70+
return {{.Ret.ReturnName}}, nil
71+
}
4772
{{end}}
4873

4974
{{if eq .Cmd ":many"}}

‎internal/codegen/golang/templates/stdlib/interfaceCode.tmpl‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,15 @@
1111
{{end -}}
1212
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
1313
{{- end}}
14+
{{- if and (eq .Cmd ":first") ($dbtxParam) }}
15+
{{range .Comments}}//{{.}}
16+
{{end -}}
17+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
18+
{{- else if eq .Cmd ":first"}}
19+
{{range .Comments}}//{{.}}
20+
{{end -}}
21+
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
22+
{{- end}}
1423
{{- if and (eq .Cmd ":many") ($dbtxParam) }}
1524
{{range .Comments}}//{{.}}
1625
{{end -}}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,26 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}
3535
{{- end}}
3636
return {{.Ret.ReturnName}}, err
3737
}
38+
{{else if eq .Cmd ":first"}}
39+
{{range .Comments}}//{{.}}
40+
{{end -}}
41+
func (q *Queries) {{.MethodName}}(ctx context.Context, {{ dbarg }} {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
42+
{{- template "queryCodeStdExec" . }}
43+
{{- if or (ne .Arg.Pair .Ret.Pair) (ne .Arg.DefineType .Ret.DefineType) }}
44+
var {{.Ret.Name}} {{.Ret.Type}}
45+
{{- end}}
46+
err := row.Scan({{.Ret.Scan}})
47+
if err != nil {
48+
if errors.Is(err, sql.ErrNoRows) {
49+
return nil, nil
50+
}
51+
{{- if $.WrapErrors}}
52+
err = fmt.Errorf("query {{.MethodName}}: %w", err)
53+
{{- end}}
54+
return nil, err
55+
}
56+
return {{.Ret.ReturnName}}, nil
57+
}
3858
{{end}}
3959

4060
{{if eq .Cmd ":many"}}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# package querytest
2-
query.sql:1:1: missing query type [':one', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
2+
query.sql:1:1: missing query type [':one', ':first', ':many', ':exec', ':execrows', ':execlastid', ':execresult', ':copyfrom', 'batchexec', 'batchmany', 'batchone']: -- name: ListFoos
33
query.sql:5:1: invalid query comment: -- name: ListFoos :one :many
44
query.sql:8:1: invalid query type: :two
55
query.sql:11:1: query "DeleteFoo" specifies parameter ":one" without containing a RETURNING clause

0 commit comments

Comments
 (0)