import code from oblast/benchmark/internal

Minimally edited to s/oblast/gsql/.

1 files changed, 89 insertions, 0 deletions

diff --git a/handle.go b/handle.go
new file mode 100644
index 0000000..7f6514d
--- /dev/null
+++ b/handle.go
@@ -0,0 +1,89 @@
+// SPDX-FileCopyrightText: 2026 Stefan Majewsky <majewsky@gmx.net>
+// SPDX-License-Identifier: Apache-2.0
+
+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/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
+}
+
+var (
+	_ Handle = &pgx.Conn{}
+	_ Handle = &pgxpool.Conn{}
+	_ Handle = pgx.Tx(&pgxpool.Tx{})
+)
+
+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))
+	}
+}
+
+var preparedStatementId atomic.Uint64
+
+type wrappedHandle struct {
+	inner Handle
+}
+
+// GSQLPrepare implements the [gsql.Handle] interface.
+func (h wrappedHandle) GSQLPrepare(ctx context.Context, query string, repeated bool) (gsql.Statement, error) {
+	if !repeated {
+		return wrappedUnpreparedStatement{query, h.inner}, nil
+	}
+
+	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()
+		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()
+	}
+}
+
+// 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 wrappedHandle) GSQLQuery(ctx context.Context, query string, args []any) (gsql.Rows, error) {
+	rows, err := h.inner.Query(ctx, query, args...)
+	return wrappedRows{rows}, err
+}