Skip to content

Commit 768ccb6

Browse files
danielmmetzdanielmmetz
authored
Add emit_methods_with_db_argument config option (#1279)
Support a new configuration parameter `emit_methods_with_db_argument `that modifies the generated method sets to provide methods like shown in WidgetInserter. If set to true, the generated *Queries omits storing a DBTX as a struct field and requires it be passed in to all method calls. In doing so, it allows callers to easily provide the connection for standalone use or for use as part of a broader transaction and makes it easy for the surrounding code to use a narrowly defined interface.
1 parent 2aa00e2 commit 768ccb6

File tree

30 files changed

+522
-96
lines changed

30 files changed

+522
-96
lines changed

‎docs/reference/config.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ packages:
1919
emit_json_tags: true
2020
emit_result_struct_pointers: false
2121
emit_params_struct_pointers: false
22+
emit_methods_with_db_argument: false
2223
json_tags_case_style: "camel"
2324
output_db_file_name: "db.go"
2425
output_models_file_name: "models.go"
@@ -57,6 +58,8 @@ Each package document has the following keys:
5758
- If true, query results are returned as pointers to structs. Queries returning multiple results are returned as slices of pointers. Defaults to `false`.
5859
- `emit_params_struct_pointers`:
5960
- If true, parameters are passed as pointers to structs. Defaults to `false`.
61+
- `emit_methods_with_db_argument`:
62+
- If true, generated methods will accept a DBTX argument instead of storing a DBTX on the `*Queries` struct. Defaults to `false`.
6063
- `json_tags_case_style`:
6164
- `camel` for camelCase, `pascal` for PascalCase, `snake` for snake_case or `none` to use the column name in the DB. Defaults to `none`.
6265
- `output_db_file_name`:

‎internal/cmd/generate.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,11 @@ func Generate(ctx context.Context, e Env, dir, filename string, stderr io.Writer
9999
return nil, err
100100
}
101101

102+
if err := config.Validate(conf); err != nil {
103+
fmt.Fprintf(stderr, "error validating %s: %s\n", base, err)
104+
return nil, err
105+
}
106+
102107
output := map[string]string{}
103108
errored := false
104109

‎internal/codegen/golang/gen.go‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,12 @@ type tmplCtx struct {
3131
// TODO: Race conditions
3232
SourceName string
3333

34-
EmitJSONTags bool
35-
EmitDBTags bool
36-
EmitPreparedQueries bool
37-
EmitInterface bool
38-
EmitEmptySlices bool
34+
EmitJSONTags bool
35+
EmitDBTags bool
36+
EmitPreparedQueries bool
37+
EmitInterface bool
38+
EmitEmptySlices bool
39+
EmitMethodsWithDBArgument bool
3940
}
4041

4142
func (t *tmplCtx) OutputQuery(sourceName string) bool {
@@ -76,18 +77,19 @@ func generate(settings config.CombinedSettings, enums []Enum, structs []Struct,
7677

7778
golang := settings.Go
7879
tctx := tmplCtx{
79-
Settings: settings.Global,
80-
EmitInterface: golang.EmitInterface,
81-
EmitJSONTags: golang.EmitJSONTags,
82-
EmitDBTags: golang.EmitDBTags,
83-
EmitPreparedQueries: golang.EmitPreparedQueries,
84-
EmitEmptySlices: golang.EmitEmptySlices,
85-
SQLPackage: SQLPackageFromString(golang.SQLPackage),
86-
Q: "`",
87-
Package: golang.Package,
88-
GoQueries: queries,
89-
Enums: enums,
90-
Structs: structs,
80+
Settings: settings.Global,
81+
EmitInterface: golang.EmitInterface,
82+
EmitJSONTags: golang.EmitJSONTags,
83+
EmitDBTags: golang.EmitDBTags,
84+
EmitPreparedQueries: golang.EmitPreparedQueries,
85+
EmitEmptySlices: golang.EmitEmptySlices,
86+
EmitMethodsWithDBArgument: golang.EmitMethodsWithDBArgument,
87+
SQLPackage: SQLPackageFromString(golang.SQLPackage),
88+
Q: "`",
89+
Package: golang.Package,
90+
GoQueries: queries,
91+
Enums: enums,
92+
Structs: structs,
9193
}
9294

9395
output := map[string]string{}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,26 @@ type DBTX interface {
66
QueryRow(context.Context, string, ...interface{}) pgx.Row
77
}
88

9+
{{ if .EmitMethodsWithDBArgument}}
10+
func New() *Queries {
11+
return &Queries{}
12+
{{- else -}}
913
func New(db DBTX) *Queries {
1014
return &Queries{db: db}
15+
{{- end}}
1116
}
1217

1318
type Queries struct {
19+
{{if not .EmitMethodsWithDBArgument}}
1420
db DBTX
21+
{{end}}
1522
}
1623

24+
{{if not .EmitMethodsWithDBArgument}}
1725
func (q *Queries) WithTx(tx pgx.Tx) *Queries {
1826
return &Queries{
1927
db: tx,
2028
}
2129
}
22-
{{end}}
30+
{{end}}
31+
{{end}}

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

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,34 @@
11
{{define "interfaceCodePgx"}}
22
type Querier interface {
3+
{{- $dbtxParam := .EmitMethodsWithDBArgument -}}
34
{{- range .GoQueries}}
4-
{{- if eq .Cmd ":one"}}
5+
{{- if and (eq .Cmd ":one") ($dbtxParam) }}
6+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
7+
{{- else if eq .Cmd ":one" }}
58
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
69
{{- end}}
7-
{{- if eq .Cmd ":many"}}
10+
{{- if and (eq .Cmd ":many") ($dbtxParam) }}
11+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error)
12+
{{- else if eq .Cmd ":many" }}
813
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error)
914
{{- end}}
10-
{{- if eq .Cmd ":exec"}}
15+
{{- if and (eq .Cmd ":exec") ($dbtxParam) }}
16+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) error
17+
{{- else if eq .Cmd ":exec" }}
1118
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) error
1219
{{- end}}
13-
{{- if eq .Cmd ":execrows"}}
20+
{{- if and (eq .Cmd ":execrows") ($dbtxParam) }}
21+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (int64, error)
22+
{{- else if eq .Cmd ":execrows" }}
1423
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, error)
1524
{{- end}}
16-
{{- if eq .Cmd ":execresult"}}
25+
{{- if and (eq .Cmd ":execresult") ($dbtxParam) }}
26+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (pgconn.CommandTag, error)
27+
{{- else if eq .Cmd ":execresult" }}
1728
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (pgconn.CommandTag, error)
1829
{{- end}}
1930
{{- end}}
2031
}
2132

2233
var _ Querier = (*Queries)(nil)
23-
{{end}}
34+
{{end}}

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

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,13 @@ type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
2222
{{if eq .Cmd ":one"}}
2323
{{range .Comments}}//{{.}}
2424
{{end -}}
25+
{{- if $.EmitMethodsWithDBArgument}}
26+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
27+
row := db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
28+
{{- else -}}
2529
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
2630
row := q.db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
31+
{{- end}}
2732
var {{.Ret.Name}} {{.Ret.Type}}
2833
err := row.Scan({{.Ret.Scan}})
2934
return {{.Ret.ReturnName}}, err
@@ -33,8 +38,13 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.De
3338
{{if eq .Cmd ":many"}}
3439
{{range .Comments}}//{{.}}
3540
{{end -}}
41+
{{- if $.EmitMethodsWithDBArgument}}
42+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {
43+
rows, err := db.Query(ctx, {{.ConstantName}}, {{.Arg.Params}})
44+
{{- else}}
3645
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {
3746
rows, err := q.db.Query(ctx, {{.ConstantName}}, {{.Arg.Params}})
47+
{{- end}}
3848
if err != nil {
3949
return nil, err
4050
}
@@ -61,17 +71,27 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.
6171
{{if eq .Cmd ":exec"}}
6272
{{range .Comments}}//{{.}}
6373
{{end -}}
74+
{{- if $.EmitMethodsWithDBArgument}}
75+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) error {
76+
_, err := db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
77+
{{- else}}
6478
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) error {
6579
_, err := q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
80+
{{- end}}
6681
return err
6782
}
6883
{{end}}
6984

7085
{{if eq .Cmd ":execrows"}}
7186
{{range .Comments}}//{{.}}
7287
{{end -}}
88+
{{if $.EmitMethodsWithDBArgument}}
89+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (int64, error) {
90+
result, err := db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
91+
{{- else -}}
7392
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, error) {
7493
result, err := q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
94+
{{- end}}
7595
if err != nil {
7696
return 0, err
7797
}
@@ -82,11 +102,16 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, er
82102
{{if eq .Cmd ":execresult"}}
83103
{{range .Comments}}//{{.}}
84104
{{end -}}
105+
{{if $.EmitMethodsWithDBArgument}}
106+
func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (pgconn.CommandTag, error) {
107+
return db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
108+
{{else}}
85109
func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (pgconn.CommandTag, error) {
86110
return q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
111+
{{- end}}
87112
}
88113
{{end}}
89114

90115
{{end}}
91116
{{end}}
92-
{{end}}
117+
{{end}}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,13 @@ type DBTX interface {
66
QueryRowContext(context.Context, string, ...interface{}) *sql.Row
77
}
88

9+
{{ if .EmitMethodsWithDBArgument}}
10+
func New() *Queries {
11+
return &Queries{}
12+
{{- else -}}
913
func New(db DBTX) *Queries {
1014
return &Queries{db: db}
15+
{{- end}}
1116
}
1217

1318
{{if .EmitPreparedQueries}}
@@ -72,7 +77,9 @@ func (q *Queries) queryRow(ctx context.Context, stmt *sql.Stmt, query string, ar
7277
{{end}}
7378

7479
type Queries struct {
80+
{{- if not .EmitMethodsWithDBArgument}}
7581
db DBTX
82+
{{- end}}
7683

7784
{{- if .EmitPreparedQueries}}
7885
tx *sql.Tx
@@ -82,6 +89,7 @@ type Queries struct {
8289
{{- end}}
8390
}
8491

92+
{{if not .EmitMethodsWithDBArgument}}
8593
func (q *Queries) WithTx(tx *sql.Tx) *Queries {
8694
return &Queries{
8795
db: tx,
@@ -93,4 +101,5 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries {
93101
{{- end}}
94102
}
95103
}
96-
{{end}}
104+
{{end}}
105+
{{end}}

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

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,34 @@
11
{{define "interfaceCodeStd"}}
22
type Querier interface {
3+
{{- $dbtxParam := .EmitMethodsWithDBArgument -}}
34
{{- range .GoQueries}}
4-
{{- if eq .Cmd ":one"}}
5+
{{- if and (eq .Cmd ":one") ($dbtxParam) }}
6+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
7+
{{- else if eq .Cmd ":one"}}
58
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error)
69
{{- end}}
7-
{{- if eq .Cmd ":many"}}
10+
{{- if and (eq .Cmd ":many") ($dbtxParam) }}
11+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error)
12+
{{- else if eq .Cmd ":many"}}
813
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error)
914
{{- end}}
10-
{{- if eq .Cmd ":exec"}}
15+
{{- if and (eq .Cmd ":exec") ($dbtxParam) }}
16+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) error
17+
{{- else if eq .Cmd ":exec"}}
1118
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) error
1219
{{- end}}
13-
{{- if eq .Cmd ":execrows"}}
20+
{{- if and (eq .Cmd ":execrows") ($dbtxParam) }}
21+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (int64, error)
22+
{{- else if eq .Cmd ":execrows"}}
1423
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, error)
1524
{{- end}}
16-
{{- if eq .Cmd ":execresult"}}
25+
{{- if and (eq .Cmd ":execresult") ($dbtxParam) }}
26+
{{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (sql.Result, error)
27+
{{- else if eq .Cmd ":execresult"}}
1728
{{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (sql.Result, error)
1829
{{- end}}
1930
{{- end}}
2031
}
2132

2233
var _ Querier = (*Queries)(nil)
23-
{{end}}
34+
{{end}}

0 commit comments

Comments
 (0)