封装 Sqlc 的 Queries 实现更方便的事务操作

Irho — Fri, 02 Aug 2024 01:46:54 +0000

SQLC 是什么

SQLC 是一个强大的开发工具，它的核心功能是将SQL查询转换成类型安全的Go代码。通过解析SQL语句和分析数据库结构，sqlc能够自动生成对应的Go结构体和函数，大大简化了数据库操作的代码编写过程。

使用sqlc，开发者可以专注于编写SQL查询，而将繁琐的Go代码生成工作交给工具完成，从而加速开发过程并提高代码质量。

SQLC 的事务实现

Sqlc 生成的代码通常包含一个Queries结构体，它封装了所有数据库操作。这个结构体实现了一个通用的Querier接口，该接口定义了所有数据库查询方法。

关键在于，sqlc生成的New函数可以接受任何实现了DBTX接口的对象，包括*sql.DB和*sql.Tx。

事务实现的核心在于利用Go的接口多态性。当你需要在事务中执行操作时，可以创建一个*sql.Tx对象，然后将其传递给New函数来创建一个新的Queries实例。这个实例会在事务的上下文中执行所有操作。

假设我们通过 pgx 连接 Postgres 数据库，并以下代码初始化 Queries。

var Pool *pgxpool.Pool
var Queries *sqlc.Queries

func init() {
    ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
    defer cancel()

    connConfig, err := pgxpool.ParseConfig("postgres://user:password@127.0.0.1:5432/db?sslmode=disable")
    if err != nil {
        panic(err)
    }

    pool, err := pgxpool.NewWithConfig(ctx, connConfig)
    if err != nil {
        panic(err)
    }
    if err := pool.Ping(ctx); err != nil {
        panic(err)
    }

    Pool = pool
    Queries = sqlc.New(pool)
}

对事务的封装

下面这段代码是一个巧妙的sqlc事务封装，它简化了在Go中使用数据库事务的过程。函数接受一个上下文和一个回调函数作为参数，这个回调函数就是用户想在事务中执行的具体操作。

func WithTransaction(ctx context.Context, callback func(qtx *sqlc.Queries) (err error)) (err error) {
    tx, err := Pool.Begin(ctx)
    if err != nil {
        return err
    }
    defer func() {
        if e := tx.Rollback(ctx); e != nil && !errors.Is(e, pgx.ErrTxClosed) {
            err = e
        }
    }()

    if err := callback(Queries.WithTx(tx)); err != nil {
        return err
    }

    return tx.Commit(ctx)
}

函数首先开始一个新的事务，然后通过延迟执行来确保事务最终会被回滚，除非它被明确提交。这是一个安全机制，防止未完成的事务占用资源。接着，函数调用用户提供的回调，传入一个带有事务上下文的查询对象，允许用户在事务中执行所需的数据库操作。

如果回调成功执行且没有错误，函数会提交事务。任何在过程中出现的错误都会导致事务回滚。这种方法既保证了数据一致性，又大大简化了错误处理。

这个封装的优雅之处在于，它将复杂的事务管理逻辑隐藏在一个简单的函数调用之后。用户可以专注于编写业务逻辑，而不必担心事务的开始、提交或回滚。

这段代码的使用方法相当直观。你可以在需要执行事务的地方调用 db.WithTransaction 函数，并传入一个函数作为参数，该函数定义了你想在事务中执行的所有数据库操作。

err := db.WithTransaction(ctx, func(qtx *sqlc.Queries) error {
    // 在这里执行你的数据库操作
    // 例如：
    _, err := qtx.CreateUser(ctx, sqlc.CreateUserParams{
        Name: "Alice",
        Email: "alice@example.com",
    })
    if err != nil {
        return err
    }

    _, err = qtx.CreatePost(ctx, sqlc.CreatePostParams{
        Title: "First Post",
        Content: "Hello, World!",
        AuthorID: newUserID,
    })
    if err != nil {
        return err
    }

    // 如果所有操作都成功，返回 nil
    return nil
})

if err != nil {
    // 处理错误
    log.Printf("transaction failed: %v", err)
} else {
    log.Println("transaction completed successfully")
}

在这个例子中，我们在事务中创建了一个用户和一个帖子。如果任何操作失败，整个事务都会回滚。如果所有操作都成功，事务会被提交。

这种方法的好处是你不需要手动管理事务的开始、提交或回滚，所有这些都由 db.WithTransaction 函数处理。你只需要专注于在事务中执行的实际数据库操作。这大大简化了代码，并减少了出错的可能性。

更进一步的封装

上面提到的这种封装方式并非毫无缺点。

这种简单的事务封装在处理嵌套事务时存在局限性。这是因为它每次都会创建一个新的事务，而不是检查是否已经在一个事务中。

为了实现嵌套事务处理，我们必须可以获得当前事务对象，但是当前事务对象是隐藏在 sqlc.Queries 内部的，所以必须我们需要扩展 sqlc.Queries。

扩展 sqlc.Queries 的结构体被我们创建为 Repositories，他扩展了 *sqlc.Queries 并添加了一个新的属性 pool，这是一个 pgxpool.Pool 类型的指针。

type Repositories struct {
    *sqlc.Queries
    pool *pgxpool.Pool
}

func NewRepositories(pool *pgxpool.Pool) *Repositories {
    return &Repositories{
        pool:    pool,
        Queries: sqlc.New(pool),
    }
}

但是当我们开始编写代码的时候就会发现，*pgxpool.Pool 并不能满足 pgx.Tx 接口，这是因为 *pgxpool.Pool 中缺少 Rollback 和 Commit 方法，他只包含用于开始事务的 Begin 方法，为了解决这个问题，我们继续扩展 Repositories 在其中添加一个新的属性 tx，并为其添加新的 NewRepositoriesTx 方法。

type Repositories struct {
    *sqlc.Queries
    tx   pgx.Tx
    pool *pgxpool.Pool
}

func NewRepositoriesTx(tx pgx.Tx) *Repositories {
    return &Repositories{
        tx:      tx,
        Queries: sqlc.New(tx),
    }
}

现在，我们的 Repositories 结构体中同时存在 pool 和 tx 属性，这可能看起来不是很优雅，为什么不能抽象出来一个统一的 TX 类型呢，其实还是上面说到的原因，即 *pgxpool.Pool 只有开始事务的方法，而没有结束事务的方法，而解决这个问题的方法之一是，再创建一个 RepositoriesTX 结构体，在其中存储 pgx.Tx 而不是 *pgxpool.Pool ，但是这样做可能又会带来新的问题，其中之一是，我们可能要为他们两者分别实现 WithTransaction 方法，至于另外一个问题，我们后面在说，现在让我们先来实现 Repositories 的 WithTransaction 方法。

func (r *Repositories) WithTransaction(ctx context.Context, fn func(qtx *Repositories) (err error)) (err error) {
    var tx pgx.Tx
    if r.tx != nil {
        tx, err = r.tx.Begin(ctx)
    } else {
        tx, err = r.pool.Begin(ctx)
    }
    if err != nil {
        return err
    }
    defer func() {
        if e := tx.Rollback(ctx); e != nil && !errors.Is(e, pgx.ErrTxClosed) {
            err = e
        }
    }()

    if err := fn(NewRepositoriesTx(tx)); err != nil {
        return err
    }

    return tx.Commit(ctx)
}

这个方法和上一章节实现的 WithTransaction 主要不同是，他是实现在 *Repositories 上面而不是全局的，这样我们就可以通过 (r *Repositories) 中的 pgx.Tx 来开始嵌套事务了。

在没有开始事务的时候，我们可以调用 repositories.WithTransaction 来开启一个新的事务。

err := db.repositories.WithTransaction(ctx, func(tx *db.Repositories) error {

    return nil
})

多级事务也是没有问题的，非常容易实现。

err := db.repositories.WithTransaction(ctx, func(tx *db.Repositories) error {
    // 假设此处进行了一些数据操作
    // 然后，开启一个嵌套事务
    return tx.WithTransaction(ctx, func(tx *db.Repositories) error {
        // 这里可以在嵌套事务中进行一些操作
        return nil
    })
})

这个封装方案有效地确保了操作的原子性，即使其中任何一个操作失败，整个事务也会被回滚，从而保障了数据的一致性。

结束语

本文介绍了一个使用 Go 和 pgx 库封装 SQLC 数据库事务的方案。

核心是 Repositories 结构体，它封装了 SQLC 查询接口和事务处理逻辑。通过 WithTransaction 方法，我们可以在现有事务上开始新的子事务或在连接池中开始新的事务，并确保在函数返回时回滚事务。

构造函数 NewRepositories 和 NewRepositoriesTx 分别用于创建普通和事务内的 Repositories 实例。

这样可以将多个数据库操作封装在一个事务中，如果任何一个操作失败，事务将被回滚，提高了代码的可维护性和可读性。

DEV Community: Irho

封装 Sqlc 的 Queries 实现更方便的事务操作

SQLC 是什么

SQLC 的事务实现

对事务的封装

更进一步的封装

结束语