69 lines
1.3 KiB
Go
69 lines
1.3 KiB
Go
package db
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
db "server/internal/db/sqlc"
|
||
|
||
"github.com/jackc/pgx/v5"
|
||
"github.com/jackc/pgx/v5/pgxpool"
|
||
)
|
||
|
||
type Store struct {
|
||
*db.Queries
|
||
pool *pgxpool.Pool
|
||
}
|
||
|
||
func NewStore(pool *pgxpool.Pool) *Store {
|
||
return &Store{
|
||
pool: pool,
|
||
Queries: db.New(pool),
|
||
}
|
||
}
|
||
|
||
func (store *Store) WithTx(ctx context.Context, fn func(q *db.Queries) error, opts ...pgx.TxOptions) error {
|
||
txOpt := pgx.TxOptions{}
|
||
if len(opts) > 0 {
|
||
txOpt = opts[0]
|
||
}
|
||
|
||
tx, err := store.pool.BeginTx(ctx, txOpt)
|
||
if err != nil {
|
||
return fmt.Errorf("begin tx: %w", err)
|
||
}
|
||
|
||
defer func() {
|
||
// 使用 context.Background() 确保即使请求 context 被取消，
|
||
// 回滚指令也能发送给 DB，从而安全地回收连接而不是销毁连接。
|
||
_ = tx.Rollback(context.Background())
|
||
}()
|
||
|
||
q := store.Queries.WithTx(tx)
|
||
|
||
if err := fn(q); err != nil {
|
||
return err // 业务侧的错误保持原样返回
|
||
}
|
||
|
||
if err := tx.Commit(ctx); err != nil {
|
||
return fmt.Errorf("commit tx: %w", err)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// WithTxResult 带返回值的事务
|
||
func WithTxResult[T any](ctx context.Context, store *Store, fn func(q *db.Queries) (T, error), opts ...pgx.TxOptions) (T, error) {
|
||
var result T
|
||
|
||
err := store.WithTx(ctx, func(q *db.Queries) error {
|
||
r, err := fn(q)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
result = r
|
||
return nil
|
||
}, opts...)
|
||
|
||
return result, err
|
||
}
|