Is there a way for sqlc to generate code that can use pgxpool

Question

I started using recently sqlc with jackc/pgx/v5. I would like to be able to use pgxpool as well, but there is really no good way to use pgxpool after sqlc takes over in the flow. For instance, this is how I initialize the connection pool:

var err error
var pool *pgxpool.Pool
if pool, err = pgxpool.New(context.Background(), url); err != nil {
    log.Panicf(...)
}
defer pool.Close()

queries := db.New(pool) // queries *Queries

Basically, I just make queries available to the entire application whenever a database interaction is needed, but I don't pass around pool.

Moreover, since sqlc manages connections automatically, I'm not sure the implications of using something like the following snippet, because there are lots of manual steps involved and some overlapping:

ctx := context.Background()

conn, _ := pool.Acquire(ctx)
tx, _ := conn.Begin(ctx)
defer tx.Rollback(ctx)

queries := db.New(pool)
queries.WithTx(tx).OneOfTheAutogeneratedQueries(ctx)

defer conn.Release()
tx.Commit(ctx)

Anyone in the same situation? Is there a better way to approach this? I was under the assumption that the auto-generated code by sqlc would provide a mechanism to manage transactions accordingly, but looks like there is still a need for the pool's reference in order to create a pgx.Tx type, in addition to close resources programmatically as well.

Answer 1

When you use sqlc it typically generates an interface that abstracts away the underlying driver-level database connection or connection pool. this interface is named DBTX by default

type DBTX interface {
    Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
    Query(context.Context, string, ...interface{}) (pgx.Rows, error)
    QueryRow(context.Context, string, ...interface{}) pgx.Row
}

func New(db DBTX) *Queries {
    return &Queries{db: db}
}

type Queries struct {
    db DBTX
}

func (q *Queries) WithTx(tx pgx.Tx) *Queries {
    return &Queries{
        db: tx,
    }
}

It's important to note a few things here

1. pgxpool.Pool and pgxpool.Conn implement the DBTX Interface
2. pgx.Tx also implements the DBTX interface
3. The queries struct takes the DBTX interface, this is to abstract the operations that you might perform on a database connection or transaction

Now if it is your wish to manually manage the database connections you can set the emit_methods_with_db_argument option to true in the sqlc.yaml file. when you do so, the generated code will look as follows

// db.go
type DBTX interface {
    Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
    Query(context.Context, string, ...interface{}) (pgx.Rows, error)
    QueryRow(context.Context, string, ...interface{}) pgx.Row
}

func New() *Queries {
    return &Queries{}
}

type Queries struct {
}

Now query methods will have the following signature

func (q *Queries) CreateAuthor(ctx context.Context, db DBTX, arg CreateAuthorParams) (Author, error) {
    row := db.QueryRow(ctx, CreateAuthor, arg.Name, arg.Bio)
    // ...
}

Notice how the generated query methods accept a DBTX as a second parameter. you can pass here a single connection, a connection pool or even a transaction

Basically, I just make queries available to the entire application whenever a database interaction is needed, but I don't pass around pool.

I am not sure if i understand your question correctly, but we can achieve this by simple dependency injection. you can define a struct that represents your API or APP like this

type App struct {
    Queries *db.Queries
    Pool    *pgxpool.Pool
    // Other fields such as Cache, Logger, etc.

}

func NewApp(pool *pgxpool.Pool) *App {
    return &App{
        Queries: db.New(),
        Pool:    pool,
    }
}

And then, in your main or wherever you're initializing the application:

func main() {
    pool, err := pgxpool.Connect(context.Background(), connectionString)
    if err != nil {
        log.Fatalf("Unable to connect to database: %v", err)
    }
    defer pool.Close()

    app := NewApp(pool)
// rest of your application
}

Now you can execute a transaction like so

func (app *App) addBookAndUpdateAuthorEarningsHandler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()

    tx, err := app.Pool.Begin(ctx)
    if err != nil {
        http.Error(w, "Failed to begin transaction", http.StatusInternalServerError)
        return
    }
    defer tx.Rollback(ctx)

    book, err = app.Queries.CreateBook(ctx, tx, newBook)
    // handle error

    err = app.Queries.UpdateAuthorEarnings(ctx, tx, updateEarningsParams)
    // handle error 

    err = tx.Commit(ctx)
    // handle error 

    w.Write([]byte("Book added and author earnings updated successfully"))
}

You will do the same for a singular query, instead of starting a transaction you will simply acquire a connection and pass that as DBTX.

That being said, you don't necessarily need to this, Exec, Query and QueryRow will acquire a connection from the connection pool and release it at the end. Without using emit_methods_with_db_argument we can do a transaction like so

package db


func (q *Queries) AddBookAndUpdateAuthorEarningsTX(ctx context.Context, author Author, book Book) error {
    // Start a transaction
    // The underlying type of db (DBTX) is *pgxpool.Pool
    tx, err := q.db.(*pgxpool.Pool).Begin(ctx)
    if err != nil {
        return err
    }
    defer tx.Rollback(ctx)
    tq := q.WithTx(tx)
    // ... (rest of the transaction remains the same)
    // (it's important to use `tq` not `q` to execute transaction queries
    err = tx.Commit(ctx)
    if err != nil {
        return err
    }

    return nil
}

Using this, your app struct needs only to contain a reference to the Queries struct

type App struct {
    Queries *db.Queries
    // Other fields such as Cache, Logger, etc.

}

func NewApp(pool *pgxpool.Pool) *App {
    return &App{
        Queries: db.New(pool),
    }
}