add full implementation

1 files changed, 137 insertions, 49 deletions

diff --git a/handle.go b/handle.go
index 7f6514d..b3017ce 100644
--- a/handle.go
+++ b/handle.go
@@ -5,85 +5,173 @@ package gg_pgx
 
 import (
 	"context"
-	"fmt"
 	"strconv"
 	"sync/atomic"
 
 	"github.com/jackc/pgx/v5"
-	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"go.xyrillian.de/gg/errext"
 	"go.xyrillian.de/gg/gsql"
 )
 
-type Handle interface {
-	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
-	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
-	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
+type connHandle struct {
+	inner *pgx.Conn
 }
 
-var (
-	_ Handle = &pgx.Conn{}
-	_ Handle = &pgxpool.Conn{}
-	_ Handle = pgx.Tx(&pgxpool.Tx{})
-)
+// GSQLPrepare implements the [gsql.Handle] interface.
+func (h connHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
+	if !repeated {
+		return wrappedUnpreparedStatement{query, h.inner}, nil
+	}
+	name := getPreparedStatementName()
+	stmt, err := h.inner.Prepare(ctx, name, query)
+	return wrappedPreparedStatement{ctx, stmt, h.inner, h.inner.Deallocate}, err
+}
 
-func Wrap(h Handle) gsql.Handle {
-	switch h := h.(type) {
-	case *pgx.Conn:
-		return wrappedHandle{h}
-	case *pgxpool.Conn:
-		return wrappedHandle{h}
-	case pgx.Tx:
-		return wrappedHandle{h}
-	default:
-		panic(fmt.Sprintf("unexpected type: %#v", h))
+// GSQLQuery implements the [gsql.Handle] interface.
+func (h connHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
+	rows, err := h.inner.Query(ctx, query, args...)
+	return wrappedRows{rows}, err
+}
+
+// GSQLClose implements the [gsql.ConnectionHandle] interface.
+func (h connHandle) GSQLClose(ctx context.Context) error {
+	return h.inner.Close(ctx)
+}
+
+// GSQLTransact implements the [gsql.ConnectionHandle] interface.
+func (h connHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
+	tx, err := h.inner.Begin(ctx)
+	if err != nil {
+		return err
 	}
+	return transact(ctx, tx, action)
 }
 
-var preparedStatementId atomic.Uint64
+////////////////////////////////////////////////////////////////////////////////
 
-type wrappedHandle struct {
-	inner Handle
+type poolHandle struct {
+	inner *pgxpool.Pool
 }
 
 // GSQLPrepare implements the [gsql.Handle] interface.
-func (h wrappedHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
+func (h poolHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
 	if !repeated {
 		return wrappedUnpreparedStatement{query, h.inner}, nil
 	}
+	name := getPreparedStatementName()
+
+	// while we have a prepared statement, we also need to hold an acquired connection
+	conn, err := h.inner.Acquire(ctx)
+	if err != nil {
+		return nil, err
+	}
+	stmt, err := conn.Conn().Prepare(ctx, name, query)
+	if err != nil {
+		conn.Release()
+		return nil, err
+	}
+	deallocate := func(ctx context.Context, name string) error {
+		err := conn.Conn().Deallocate(ctx, name)
+		conn.Release()
+		return err
+	}
+	return wrappedPreparedStatement{ctx, stmt, h.inner, deallocate}, nil
+}
+
+// GSQLQuery implements the [gsql.Handle] interface.
+func (h poolHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
+	rows, err := h.inner.Query(ctx, query, args...)
+	return wrappedRows{rows}, err
+}
+
+// GSQLClose implements the [gsql.ConnectionHandle] interface.
+func (h poolHandle) GSQLClose(ctx context.Context) error {
+	h.inner.Close()
+	return nil
+}
+
+// GSQLTransact implements the [gsql.ConnectionHandle] interface.
+func (h poolHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
+	tx, err := h.inner.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	return transact(ctx, tx, action)
+}
 
-	name := "oblast_pgx_" + strconv.FormatUint(preparedStatementId.Add(1), 10)
-	switch inner := h.inner.(type) {
-	case *pgx.Conn:
-		stmt, err := inner.Prepare(ctx, name, query)
-		return wrappedPreparedStatement{ctx, stmt, h.inner}, err
-	case *pgxpool.Conn:
-		// pgxpool.Conn does not have Prepare()
+////////////////////////////////////////////////////////////////////////////////
+
+type poolConnHandle struct {
+	inner *pgxpool.Conn
+}
+
+// GSQLPrepare implements the [gsql.Handle] interface.
+func (h poolConnHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
+	if !repeated {
 		return wrappedUnpreparedStatement{query, h.inner}, nil
-	case pgx.Tx:
-		stmt, err := inner.Conn().Prepare(ctx, name, query)
-		return wrappedPreparedStatement{ctx, stmt, h.inner}, err
-	default:
-		panic("unreachable") // because of the check in func Wrap()
 	}
+	name := getPreparedStatementName()
+	stmt, err := h.inner.Conn().Prepare(ctx, name, query)
+	return wrappedPreparedStatement{ctx, stmt, h.inner, h.inner.Conn().Deallocate}, err
 }
 
-// Releases a prepared statement.
-func deallocate(ctx context.Context, h Handle, stmt *pgconn.StatementDescription) error {
-	switch h := h.(type) {
-	case *pgx.Conn:
-		return h.Deallocate(ctx, stmt.Name)
-	case *pgxpool.Conn:
-		panic("unreachable") // because func GSQLPrepare() does not return a wrappedPreparedStatement for this underlying type
-	case pgx.Tx:
-		return h.Conn().Deallocate(ctx, stmt.Name)
-	default:
-		panic("unreachable") // because of the check in func Wrap()
+// GSQLQuery implements the [gsql.Handle] interface.
+func (h poolConnHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
+	rows, err := h.inner.Query(ctx, query, args...)
+	return wrappedRows{rows}, err
+}
+
+// GSQLClose implements the [gsql.ConnectionHandle] interface.
+func (h poolConnHandle) GSQLClose(ctx context.Context) error {
+	h.inner.Release()
+	return nil
+}
+
+// GSQLTransact implements the [gsql.ConnectionHandle] interface.
+func (h poolConnHandle) GSQLTransact(ctx context.Context, action func(tx gsql.Handle) error) error {
+	tx, err := h.inner.Begin(ctx)
+	if err != nil {
+		return err
 	}
+	return transact(ctx, tx, action)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+type txHandle struct {
+	inner pgx.Tx
+}
+
+// GSQLPrepare implements the [gsql.Handle] interface.
+func (h txHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
+	if !repeated {
+		return wrappedUnpreparedStatement{query, h.inner}, nil
+	}
+	name := getPreparedStatementName()
+	stmt, err := h.inner.Conn().Prepare(ctx, name, query)
+	return wrappedPreparedStatement{ctx, stmt, h.inner, h.inner.Conn().Deallocate}, err
 }
 
 // GSQLQuery implements the [gsql.Handle] interface.
-func (h wrappedHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
+func (h txHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
 	rows, err := h.inner.Query(ctx, query, args...)
 	return wrappedRows{rows}, err
 }
+
+////////////////////////////////////////////////////////////////////////////////
+
+var preparedStatementId atomic.Uint64
+
+func getPreparedStatementName() string {
+	return "oblast_pgx_" + strconv.FormatUint(preparedStatementId.Add(1), 10)
+}
+
+func transact(ctx context.Context, tx pgx.Tx, action func(tx gsql.Handle) error) error {
+	err := action(NewTx(tx))
+	if err == nil {
+		return errext.WithCleanup(nil, "tx.Commit", tx.Commit(ctx))
+	} else {
+		return errext.WithCleanup(err, "tx.Rollback", tx.Rollback(ctx))
+	}
+}