有没有办法让sqlc生成可以使用pgxpool的代码

Question

sqlc我最近开始使用jackc/pgx/v5. 我也希望能够使用，但是在流程接管后pgxpool确实没有什么好的方法可以使用。例如，这就是我初始化连接池的方式：pgxpoolsqlc

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

基本上，queries每当需要数据库交互时，我只是将其提供给整个应用程序，但我不会传递pool.

此外，由于sqlc自动管理连接，我不确定使用类似以下代码片段的含义，因为涉及很多手动步骤并且有些重叠：

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)

有人有同样的情况吗？有更好的方法来解决这个问题吗？我假设自动生成的代码将提供一种相应的管理事务的机制，但看起来除了以编程方式关闭资源之外，sqlc仍然需要池的引用才能创建类型。pgx.Tx

Answer 1

当您使用sqlc它时，它通常会生成一个接口，该接口抽象出底层驱动程序级数据库连接或连接池。该接口被DBTX默认命名

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,
    }
}

这里需要注意一些事情

1. pgxpool.Pool并pgxpool.Conn实现DBTX接口
2. pgx.Tx还实现了该DBTX接口
3. 查询结构采用DBTX接口，这是为了抽象您可能在数据库连接或事务上执行的操作

现在，如果您希望手动管理数据库连接，您可以emit_methods_with_db_argument在文件中将该选项设置为 true sqlc.yaml。当您这样做时，生成的代码将如下所示

// 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 {
}

现在查询方法将具有以下签名

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

请注意生成的查询方法如何接受 aDBTX作为第二个参数。您可以在此处传递单个连接、连接池甚至事务

基本上，queries每当需要数据库交互时，我只是将其提供给整个应用程序，但我不会传递池。

我不确定我是否正确理解你的问题，但我们可以通过简单的依赖注入来实现这一点。你可以像这样定义一个代表你的API或APP的结构体

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,
    }
}

然后，在您的主目录或您正在初始化应用程序的任何地方：

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
}

现在您可以像这样执行交易

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"))
}

您将对单个查询执行相同的操作，您只需获取一个连接并将其作为 传递，而不是启动事务DBTX。

话虽如此，您不一定需要这样做，并且Exec将从连接池中获取连接并在最后释放它。不使用我们可以像这样进行交易QueryQueryRowemit_methods_with_db_argument

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
}

Queries使用它，您的应用程序结构只需要包含对该结构的引用

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),
    }
}