Extended postgres interface

Author: xdoubleu

pkg/database/postgres/interface.go

@@ -21,6 +21,8 @@ type DB interface {
 		optionsAndArgs ...any,
 	) (pgx.Rows, error)
 	QueryRow(ctx context.Context, sql string, optionsAndArgs ...any) pgx.Row
+	SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults
 	Begin(ctx context.Context) (pgx.Tx, error)
+	BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error)
 	Ping(ctx context.Context) error
 }

pkg/database/postgres/pgx_sync_tx.go

@@ -106,6 +106,17 @@ func (tx *PgxSyncTx) Query(
 	)
 }
 
+// SendBatch is used to wrap [pgx.Tx.QueryRow] in a [database.SyncTx].
+func (tx *PgxSyncTx) SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults {
+	return database.WrapInSyncTxNoError(
+		ctx,
+		tx.syncTx,
+		func(ctx context.Context) pgx.BatchResults {
+			return tx.syncTx.Tx.SendBatch(ctx, b)
+		},
+	)
+}
+
 // Close doesn't do anything for [PgxSyncRows] as these are closed in [Query].
 func (rows *PgxSyncRows) Close() {
 }
@@ -208,6 +219,28 @@ func (tx *PgxSyncTx) Begin(ctx context.Context) (pgx.Tx, error) {
 	)
 }
 
+// BeginTx is used to wrap [pgx.Tx.BeginTx] in a [database.SyncTx].
+func (tx *PgxSyncTx) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) {
+	return database.WrapInSyncTx(
+		ctx,
+		tx.syncTx,
+		func(ctx context.Context) (pgx.Tx, error) {
+			return tx.syncTx.Tx.Conn().BeginTx(ctx, txOptions)
+		},
+	)
+}
+
+// Commit is used to wrap [pgx.Tx.Commit] in a [database.SyncTx].
+func (tx *PgxSyncTx) Commit(ctx context.Context) error {
+	return database.WrapInSyncTxNoError(
+		ctx,
+		tx.syncTx,
+		func(ctx context.Context) error {
+			return tx.syncTx.Tx.Commit(ctx)
+		},
+	)
+}
+
 // Rollback is used to wrap [pgx.Tx.Rollback] in a [database.SyncTx].
 func (tx *PgxSyncTx) Rollback(ctx context.Context) error {
 	return database.WrapInSyncTxNoError(

pkg/database/postgres/spandb.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/XDoubleU/essentia/pkg/database"
 	"github.com/jackc/pgx/v5"
@@ -54,12 +55,34 @@ func (db *SpanDB) QueryRow(
 		optionsAndArgs...)
 }
 
+// SendBatch is used to wrap SendBatch in a [sentry.Span].
+func (db *SpanDB) SendBatch(
+	ctx context.Context,
+	b *pgx.Batch,
+) pgx.BatchResults {
+	sql := ""
+	for i, query := range b.QueuedQueries {
+		sql += fmt.Sprintf("query %d: %s\n", i, query.SQL)
+	}
+
+	span := database.StartSpan(ctx, db.dbName, sql)
+	defer span.Finish()
+
+	return db.DB.SendBatch(ctx, b)
+}
+
 // Begin doesn't wrap Begin in a [sentry.Span] as
 // this makes little sense for starting a transaction.
 func (db *SpanDB) Begin(ctx context.Context) (pgx.Tx, error) {
 	return db.DB.Begin(ctx)
 }
 
+// BeginTx doesn't wrap BeginTx in a [sentry.Span] as
+// this makes little sense for starting a transaction.
+func (db *SpanDB) BeginTx(ctx context.Context, txOptions pgx.TxOptions) (pgx.Tx, error) {
+	return db.DB.BeginTx(ctx, txOptions)
+}
+
 // Ping doesn't wrap Ping in a [sentry.Span] as
 // this makes little sense for pinging the db.
 func (db *SpanDB) Ping(ctx context.Context) error {

pkg/database/span.go

@@ -6,7 +6,7 @@ import (
 	"github.com/getsentry/sentry-go"
 )
 
-func startSpan(ctx context.Context, dbName string, sql string) *sentry.Span {
+func StartSpan(ctx context.Context, dbName string, sql string) *sentry.Span {
 	span := sentry.StartSpan(ctx, "db.query", sentry.WithDescription(sql))
 	span.SetData("db.system", dbName)
 
@@ -20,7 +20,7 @@ func WrapWithSpan[T any](
 	dbName string,
 	queryFunc func(ctx context.Context, sql string, args ...any) (T, error),
 	sql string, args ...any) (T, error) {
-	span := startSpan(ctx, dbName, sql)
+	span := StartSpan(ctx, dbName, sql)
 	defer span.Finish()
 
 	return queryFunc(ctx, sql, args...)
@@ -34,7 +34,7 @@ func WrapWithSpanNoError[T any](
 	dbName string,
 	queryFunc func(ctx context.Context, sql string, args ...any) T,
 	sql string, args ...any) T {
-	span := startSpan(ctx, dbName, sql)
+	span := StartSpan(ctx, dbName, sql)
 	defer span.Finish()
 
 	return queryFunc(ctx, sql, args...)