feat(codegen): expose DBTX inside Queries when expose_db_connection is enabled

Signed-off-by: Nathanael DEMACON <[email protected]>

Author: quantumsheep

docs/reference/config.md

@@ -189,6 +189,8 @@ The `gen` mapping supports the following keys:
   - Customize the name of the copyfrom file. Defaults to `copyfrom.go`.
 - `output_files_suffix`:
   - If specified the suffix will be added to the name of the generated files.
+- `expose_db_connection`:
+  - If true, expose the `DB` field on the `Queries` struct. Defaults to `false`.
 - `query_parameter_limit`:
   - The number of positional arguments that will be generated for Go functions. To always emit a parameter struct, set this to `0`. Defaults to `1`.
 - `rename`:

internal/codegen/golang/gen.go

@@ -41,6 +41,7 @@ type tmplCtx struct {
 	UsesBatch                 bool
 	OmitSqlcVersion           bool
 	BuildTags                 string
+	ExposeDbConnection        bool
 }
 
 func (t *tmplCtx) OutputQuery(sourceName string) bool {
@@ -187,6 +188,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 		SqlcVersion:               req.SqlcVersion,
 		BuildTags:                 options.BuildTags,
 		OmitSqlcVersion:           options.OmitSqlcVersion,
+		ExposeDbConnection:        options.ExposeDbConnection,
 	}
 
 	if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() && options.SqlDriver != opts.SQLDriverGoSQLDriverMySQL {

internal/codegen/golang/opts/options.go

@@ -38,6 +38,7 @@ type Options struct {
 	OutputQuerierFileName       string            `json:"output_querier_file_name,omitempty" yaml:"output_querier_file_name"`
 	OutputCopyfromFileName      string            `json:"output_copyfrom_file_name,omitempty" yaml:"output_copyfrom_file_name"`
 	OutputFilesSuffix           string            `json:"output_files_suffix,omitempty" yaml:"output_files_suffix"`
+	ExposeDbConnection          bool              `json:"expose_db_connection,omitempty" yaml:"expose_db_connection"`
 	InflectionExcludeTableNames []string          `json:"inflection_exclude_table_names,omitempty" yaml:"inflection_exclude_table_names"`
 	QueryParameterLimit         *int32            `json:"query_parameter_limit,omitempty" yaml:"query_parameter_limit"`
 	OmitSqlcVersion             bool              `json:"omit_sqlc_version,omitempty" yaml:"omit_sqlc_version"`
@@ -150,6 +151,9 @@ func ValidateOpts(opts *Options) error {
 	if *opts.QueryParameterLimit < 0 {
 		return fmt.Errorf("invalid options: query parameter limit must not be negative")
 	}
+	if opts.ExposeDbConnection && opts.EmitMethodsWithDbArgument {
+		return fmt.Errorf("invalid options: expose_db_connection and emit_methods_with_db_argument options are mutually exclusive")
+	}
 
 	return nil
 }

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

@@ -34,4 +34,10 @@ func (q *Queries) WithTx(tx pgx.Tx) *Queries {
 	}
 }
 {{end}}
+
+{{if .ExposeDbConnection}}
+func (q *Queries) Conn() DBTX {
+	return q.db
+}
+{{end}}
 {{end}}

internal/config/v_two.json

@@ -143,6 +143,9 @@
                                     "emit_sql_as_comment": {
                                         "type": "boolean"
                                     },
+                                    "expose_db_connection": {
+                                        "type": "boolean"
+                                    },
                                     "build_tags": {
                                         "type": "string"
                                     },

internal/endtoend/testdata/expose_db_connection/db/db.go

@@ -0,0 +1,19 @@
+-- name: GetAuthor :one
+SELECT * FROM authors
+WHERE id = $1 LIMIT 1;
+
+-- name: ListAuthors :many
+SELECT * FROM authors
+ORDER BY name;
+
+-- name: CreateAuthor :one
+INSERT INTO authors (
+  name, bio
+) VALUES (
+  $1, $2
+)
+RETURNING *;
+
+-- name: DeleteAuthor :exec
+DELETE FROM authors
+WHERE id = $1;

internal/endtoend/testdata/expose_db_connection/db/models.go

@@ -0,0 +1,5 @@
+CREATE TABLE authors (
+  id   BIGSERIAL PRIMARY KEY,
+  name text      NOT NULL,
+  bio  text
+);