1 files changed, 31 insertions, 16 deletions
diff --git a/statement.go b/statement.go
index bb95252..1cda044 100644
--- a/statement.go
+++ b/statement.go
@@ -7,19 +7,34 @@ import (
 	"context"
 	"database/sql"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"go.xyrillian.de/gg/gsql"
 )
 
+type pgxExecutor interface {
+	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
+	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) // TODO: remove after splitting Handle types
+	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
+}
+
+var (
+	_ pgxExecutor = &pgx.Conn{}
+	_ pgxExecutor = &pgxpool.Conn{}
+	_ pgxExecutor = pgx.Tx(&pgxpool.Tx{})
+)
+
 type wrappedPreparedStatement struct {
-	ctx       context.Context
-	statement *pgconn.StatementDescription
-	handle    Handle
+	ctx        context.Context
+	statement  *pgconn.StatementDescription
+	executor   pgxExecutor
+	deallocate func(ctx context.Context, name string) error
 }
 
 type wrappedUnpreparedStatement struct {
-	query  string
-	handle Handle
+	query    string
+	executor pgxExecutor
 }
 
 var (
@@ -27,34 +42,34 @@ var (
 	_ gsql.Statement = wrappedUnpreparedStatement{}
 )
 
-// Close implements the [handle.Statement] interface.
+// Close implements the [gsql.Statement] interface.
 func (s wrappedPreparedStatement) Close() error {
-	return deallocate(s.ctx, s.handle, s.statement)
+	return s.deallocate(s.ctx, s.statement.Name)
 }
 
-// Close implements the [handle.Statement] interface.
+// Close implements the [gsql.Statement] interface.
 func (s wrappedUnpreparedStatement) Close() error {
 	return nil
 }
 
-// Exec implements the [handle.Statement] interface.
+// Exec implements the [gsql.Statement] interface.
 func (s wrappedPreparedStatement) Exec(ctx context.Context, args []any) (sql.Result, error) {
-	result, err := s.handle.Exec(ctx, s.statement.Name, args...)
+	result, err := s.executor.Exec(ctx, s.statement.Name, args...)
 	return wrappedResult{result}, err
 }
 
-// Exec implements the [handle.Statement] interface.
+// Exec implements the [gsql.Statement] interface.
 func (s wrappedUnpreparedStatement) Exec(ctx context.Context, args []any) (sql.Result, error) {
-	result, err := s.handle.Exec(ctx, s.query, args...)
+	result, err := s.executor.Exec(ctx, s.query, args...)
 	return wrappedResult{result}, err
 }
 
-// QueryRow implements the [handle.Statement] interface.
+// QueryRow implements the [gsql.Statement] interface.
 func (s wrappedPreparedStatement) QueryRow(ctx context.Context, args, slots []any) error {
-	return s.handle.QueryRow(ctx, s.statement.Name, args...).Scan(slots...)
+	return s.executor.QueryRow(ctx, s.statement.Name, args...).Scan(slots...)
 }
 
-// QueryRow implements the [handle.Statement] interface.
+// QueryRow implements the [gsql.Statement] interface.
 func (s wrappedUnpreparedStatement) QueryRow(ctx context.Context, args, slots []any) error {
-	return s.handle.QueryRow(ctx, s.query, args...).Scan(slots...)
+	return s.executor.QueryRow(ctx, s.query, args...).Scan(slots...)
 }