From 38ffa181e18d4cebbd7c387ffc506848a5def269 Mon Sep 17 00:00:00 2001
From: Michael Li <alimy@gility.net>
Date: Sat, 7 Jan 2023 19:06:17 +0800
Subject: [PATCH] sqlc: base pgx servant add custom context help method

---
 internal/dao/slonik/pgx.go | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/internal/dao/slonik/pgx.go b/internal/dao/slonik/pgx.go
index d14c5f07..3558df0a 100644
--- a/internal/dao/slonik/pgx.go
+++ b/internal/dao/slonik/pgx.go
@@ -63,6 +63,28 @@ func (s *pgxServant) withTx(txOptions pgx.TxOptions, handle func(dbr.Querier) er
 	return tx.Rollback(ctx)
 }
 
+func (s *pgxServant) withCtx(ctx context.Context, handle func(dbr.Querier) error) error {
+	tx, err := s.db.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	if err = handle(dbr.New(tx)); err == nil {
+		return tx.Commit(ctx)
+	}
+	return tx.Rollback(ctx)
+}
+
+func (s *pgxServant) withTxCtx(ctx context.Context, txOptions pgx.TxOptions, handle func(dbr.Querier) error) error {
+	tx, err := s.db.BeginTx(ctx, txOptions)
+	if err != nil {
+		return err
+	}
+	if err = handle(dbr.New(tx)); err == nil {
+		return tx.Commit(ctx)
+	}
+	return tx.Rollback(ctx)
+}
+
 func newPgxServant(db *pgx.Conn) *pgxServant {
 	return &pgxServant{
 		db: db,