From 4808ce2509ad6bb13c742daa63fa41fa6fa337d1 Mon Sep 17 00:00:00 2001
From: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
Date: Thu, 23 Jan 2025 12:57:35 +0100
Subject: [PATCH 1/2] feat(codegen): expose DBTX inside Queries when
`expose_db_connection` is enabled
Signed-off-by: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
---
docs/reference/config.md | 2 +
internal/codegen/golang/gen.go | 2 +
internal/codegen/golang/opts/options.go | 4 +
.../codegen/golang/templates/pgx/dbCode.tmpl | 6 ++
internal/config/v_two.json | 3 +
.../testdata/expose_db_connection/db/db.go | 35 ++++++++
.../expose_db_connection/db/models.go | 15 ++++
.../expose_db_connection/db/query.sql.go | 82 +++++++++++++++++++
.../testdata/expose_db_connection/query.sql | 19 +++++
.../testdata/expose_db_connection/schema.sql | 5 ++
.../testdata/expose_db_connection/sqlc.json | 15 ++++
11 files changed, 188 insertions(+)
create mode 100644 internal/endtoend/testdata/expose_db_connection/db/db.go
create mode 100644 internal/endtoend/testdata/expose_db_connection/db/models.go
create mode 100644 internal/endtoend/testdata/expose_db_connection/db/query.sql.go
create mode 100644 internal/endtoend/testdata/expose_db_connection/query.sql
create mode 100644 internal/endtoend/testdata/expose_db_connection/schema.sql
create mode 100644 internal/endtoend/testdata/expose_db_connection/sqlc.json
diff --git a/docs/reference/config.md b/docs/reference/config.md
index 9d334834f0..eef2694df4 100644
--- a/docs/reference/config.md
+++ b/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`:
diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
index 5b7977f500..be3f32766e 100644
--- a/internal/codegen/golang/gen.go
+++ b/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 {
diff --git a/internal/codegen/golang/opts/options.go b/internal/codegen/golang/opts/options.go
index 30a6c2246c..9b321880cc 100644
--- a/internal/codegen/golang/opts/options.go
+++ b/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
}
diff --git a/internal/codegen/golang/templates/pgx/dbCode.tmpl b/internal/codegen/golang/templates/pgx/dbCode.tmpl
index 236554d9f2..6fec262258 100644
--- a/internal/codegen/golang/templates/pgx/dbCode.tmpl
+++ b/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}}
diff --git a/internal/config/v_two.json b/internal/config/v_two.json
index acf914997d..19ffd42c84 100644
--- a/internal/config/v_two.json
+++ b/internal/config/v_two.json
@@ -143,6 +143,9 @@
"emit_sql_as_comment": {
"type": "boolean"
},
+ "expose_db_connection": {
+ "type": "boolean"
+ },
"build_tags": {
"type": "string"
},
diff --git a/internal/endtoend/testdata/expose_db_connection/db/db.go b/internal/endtoend/testdata/expose_db_connection/db/db.go
new file mode 100644
index 0000000000..c83c3a95bb
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/db/db.go
@@ -0,0 +1,35 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+// sqlc v1.28.0
+
+package db
+
+import (
+ "context"
+ "database/sql"
+)
+
+type DBTX interface {
+ ExecContext(context.Context, string, ...interface{}) (sql.Result, error)
+ PrepareContext(context.Context, string) (*sql.Stmt, error)
+ QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error)
+ QueryRowContext(context.Context, string, ...interface{}) *sql.Row
+}
+
+func New(db DBTX) *Queries {
+ return &Queries{db: db}
+}
+
+type Queries struct {
+ db DBTX
+}
+
+func (q *Queries) WithTx(tx *sql.Tx) *Queries {
+ return &Queries{
+ db: tx,
+ }
+}
+
+func (q *Queries) Conn() DBTX {
+ return q.db
+}
diff --git a/internal/endtoend/testdata/expose_db_connection/db/models.go b/internal/endtoend/testdata/expose_db_connection/db/models.go
new file mode 100644
index 0000000000..83aad4e56d
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/db/models.go
@@ -0,0 +1,15 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+// sqlc v1.28.0
+
+package db
+
+import (
+ "database/sql"
+)
+
+type Author struct {
+ ID int64
+ Name string
+ Bio sql.NullString
+}
diff --git a/internal/endtoend/testdata/expose_db_connection/db/query.sql.go b/internal/endtoend/testdata/expose_db_connection/db/query.sql.go
new file mode 100644
index 0000000000..4731e891ce
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/db/query.sql.go
@@ -0,0 +1,82 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+// sqlc v1.28.0
+// source: query.sql
+
+package db
+
+import (
+ "context"
+ "database/sql"
+)
+
+const createAuthor = `-- name: CreateAuthor :one
+INSERT INTO authors (
+ name, bio
+) VALUES (
+ $1, $2
+)
+RETURNING id, name, bio
+`
+
+type CreateAuthorParams struct {
+ Name string
+ Bio sql.NullString
+}
+
+func (q *Queries) CreateAuthor(ctx context.Context, arg CreateAuthorParams) (Author, error) {
+ row := q.db.QueryRowContext(ctx, createAuthor, arg.Name, arg.Bio)
+ var i Author
+ err := row.Scan(&i.ID, &i.Name, &i.Bio)
+ return i, err
+}
+
+const deleteAuthor = `-- name: DeleteAuthor :exec
+DELETE FROM authors
+WHERE id = $1
+`
+
+func (q *Queries) DeleteAuthor(ctx context.Context, id int64) error {
+ _, err := q.db.ExecContext(ctx, deleteAuthor, id)
+ return err
+}
+
+const getAuthor = `-- name: GetAuthor :one
+SELECT id, name, bio FROM authors
+WHERE id = $1 LIMIT 1
+`
+
+func (q *Queries) GetAuthor(ctx context.Context, id int64) (Author, error) {
+ row := q.db.QueryRowContext(ctx, getAuthor, id)
+ var i Author
+ err := row.Scan(&i.ID, &i.Name, &i.Bio)
+ return i, err
+}
+
+const listAuthors = `-- name: ListAuthors :many
+SELECT id, name, bio FROM authors
+ORDER BY name
+`
+
+func (q *Queries) ListAuthors(ctx context.Context) ([]Author, error) {
+ rows, err := q.db.QueryContext(ctx, listAuthors)
+ if err != nil {
+ return nil, err
+ }
+ defer rows.Close()
+ var items []Author
+ for rows.Next() {
+ var i Author
+ if err := rows.Scan(&i.ID, &i.Name, &i.Bio); err != nil {
+ return nil, err
+ }
+ items = append(items, i)
+ }
+ if err := rows.Close(); err != nil {
+ return nil, err
+ }
+ if err := rows.Err(); err != nil {
+ return nil, err
+ }
+ return items, nil
+}
diff --git a/internal/endtoend/testdata/expose_db_connection/query.sql b/internal/endtoend/testdata/expose_db_connection/query.sql
new file mode 100644
index 0000000000..971b8f902d
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/query.sql
@@ -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;
diff --git a/internal/endtoend/testdata/expose_db_connection/schema.sql b/internal/endtoend/testdata/expose_db_connection/schema.sql
new file mode 100644
index 0000000000..69b607d902
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/schema.sql
@@ -0,0 +1,5 @@
+CREATE TABLE authors (
+ id BIGSERIAL PRIMARY KEY,
+ name text NOT NULL,
+ bio text
+);
diff --git a/internal/endtoend/testdata/expose_db_connection/sqlc.json b/internal/endtoend/testdata/expose_db_connection/sqlc.json
new file mode 100644
index 0000000000..71b4863840
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/sqlc.json
@@ -0,0 +1,15 @@
+{
+ "version": "2",
+ "sql": [{
+ "schema": "schema.sql",
+ "queries": "query.sql",
+ "engine": "postgresql",
+ "gen": {
+ "go": {
+ "out": "db",
+ "expose_db_connection": true
+ }
+ }
+ }]
+}
+
From 299907e097644ae350e2463b3003378332cfc3ab Mon Sep 17 00:00:00 2001
From: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
Date: Thu, 23 Jan 2025 14:50:15 +0100
Subject: [PATCH 2/2] add exec.json
Signed-off-by: Nathanael DEMACON <quantumsheep@users.noreply.github.com>
---
.../testdata/expose_db_connection/exec.json | 3 +++
.../testdata/expose_db_connection/sqlc.json | 15 ---------------
.../testdata/expose_db_connection/sqlc.yaml | 9 +++++++++
3 files changed, 12 insertions(+), 15 deletions(-)
create mode 100644 internal/endtoend/testdata/expose_db_connection/exec.json
delete mode 100644 internal/endtoend/testdata/expose_db_connection/sqlc.json
create mode 100644 internal/endtoend/testdata/expose_db_connection/sqlc.yaml
diff --git a/internal/endtoend/testdata/expose_db_connection/exec.json b/internal/endtoend/testdata/expose_db_connection/exec.json
new file mode 100644
index 0000000000..ee1b7ecd9e
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/exec.json
@@ -0,0 +1,3 @@
+{
+ "contexts": ["managed-db"]
+}
diff --git a/internal/endtoend/testdata/expose_db_connection/sqlc.json b/internal/endtoend/testdata/expose_db_connection/sqlc.json
deleted file mode 100644
index 71b4863840..0000000000
--- a/internal/endtoend/testdata/expose_db_connection/sqlc.json
+++ /dev/null
@@ -1,15 +0,0 @@
-{
- "version": "2",
- "sql": [{
- "schema": "schema.sql",
- "queries": "query.sql",
- "engine": "postgresql",
- "gen": {
- "go": {
- "out": "db",
- "expose_db_connection": true
- }
- }
- }]
-}
-
diff --git a/internal/endtoend/testdata/expose_db_connection/sqlc.yaml b/internal/endtoend/testdata/expose_db_connection/sqlc.yaml
new file mode 100644
index 0000000000..292af899db
--- /dev/null
+++ b/internal/endtoend/testdata/expose_db_connection/sqlc.yaml
@@ -0,0 +1,9 @@
+version: "2"
+sql:
+ - engine: "postgresql"
+ schema: "schema.sql"
+ queries: "query.sql"
+ gen:
+ go:
+ out: "db"
+ expose_db_connection: true