driver.go

package main

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"io"
	"sync"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Driver implements driver.Driver.
type Driver struct{}

// Open returns a new connection to the database.
func (d *Driver) Open(name string) (driver.Conn, error) {
	return nil, errors.New("use OpenDB with Connector")
}

// Connector implements driver.Connector.
type Connector struct {
	pool *pgxpool.Pool
}

// NewConnector creates a new Connector wrapping the given pgxpool.Pool.
func NewConnector(pool *pgxpool.Pool) *Connector {
	return &Connector{pool: pool}
}

// Connect returns a connection to the database.
func (c *Connector) Connect(ctx context.Context) (driver.Conn, error) {
	conn, err := c.pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	return &Conn{conn: conn}, nil
}

// Driver returns the underlying Driver.
func (c *Connector) Driver() driver.Driver {
	return &Driver{}
}

// Conn implements driver.Conn, driver.ConnBeginTx, driver.QueryerContext, and driver.ExecerContext.
type Conn struct {
	conn *pgxpool.Conn
	bad  bool
	mu   sync.Mutex
}

// Prepare returns a prepared statement, bound to this connection.
func (c *Conn) Prepare(query string) (driver.Stmt, error) {
	return nil, errors.New("prepared statements not implemented")
}

// Close invalidates and returns the connection to the pool.
func (c *Conn) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.bad {
		// Close the underlying connection so pgxpool discards it
		c.conn.Conn().Close(context.Background())
	}
	c.conn.Release()
	return nil
}

// Begin starts a transaction. Deprecated: use BeginTx.
func (c *Conn) Begin() (driver.Tx, error) {
	return nil, errors.New("use BeginTx")
}

// BeginTx starts a transaction with context and options.
func (c *Conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.bad {
		return nil, driver.ErrBadConn
	}

	pgxOpts := pgx.TxOptions{}
	if opts.ReadOnly {
		pgxOpts.AccessMode = pgx.ReadOnly
	}
	switch sql.IsolationLevel(opts.Isolation) {
	case sql.LevelDefault:
	case sql.LevelReadUncommitted:
		pgxOpts.IsoLevel = pgx.ReadUncommitted
	case sql.LevelReadCommitted:
		pgxOpts.IsoLevel = pgx.ReadCommitted
	case sql.LevelRepeatableRead:
		pgxOpts.IsoLevel = pgx.RepeatableRead
	case sql.LevelSerializable:
		pgxOpts.IsoLevel = pgx.Serializable
	}

	tx, err := c.conn.BeginTx(ctx, pgxOpts)
	if err != nil {
		return nil, err
	}
	return &Tx{tx: tx, conn: c, ctx: ctx}, nil
}

// QueryContext executes a query that returns rows.
func (c *Conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.bad {
		return nil, driver.ErrBadConn
	}
	pgxArgs := make([]any, len(args))
	for i, arg := range args {
		pgxArgs[i] = arg.Value
	}
	rows, err := c.conn.Query(ctx, query, pgxArgs...)
	if err != nil {
		return nil, err
	}
	return &Rows{rows: rows}, nil
}

// ExecContext executes a query that doesn't return rows.
func (c *Conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.bad {
		return nil, driver.ErrBadConn
	}
	pgxArgs := make([]any, len(args))
	for i, arg := range args {
		pgxArgs[i] = arg.Value
	}
	tag, err := c.conn.Exec(ctx, query, pgxArgs...)
	if err != nil {
		return nil, err
	}
	return driver.RowsAffected(tag.RowsAffected()), nil
}

// Tx implements driver.Tx.
type Tx struct {
	tx   pgx.Tx
	conn *Conn
	ctx  context.Context
}

// Commit commits the transaction.
func (t *Tx) Commit() error {
	err := t.tx.Commit(t.ctx)
	if err != nil {
		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
			t.conn.mu.Lock()
			t.conn.bad = true
			t.conn.mu.Unlock()
			return driver.ErrBadConn
		}
		return err
	}
	return nil
}

// Rollback rollbacks the transaction.
func (t *Tx) Rollback() error {
	err := t.tx.Rollback(t.ctx)
	if err != nil {
		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
			t.conn.mu.Lock()
			t.conn.bad = true
			t.conn.mu.Unlock()
			return driver.ErrBadConn
		}
		return err
	}
	return nil
}

// Rows implements driver.Rows.
type Rows struct {
	rows pgx.Rows
}

// Columns returns the names of the columns.
func (r *Rows) Columns() []string {
	fields := r.rows.FieldDescriptions()
	cols := make([]string, len(fields))
	for i, f := range fields {
		cols[i] = f.Name
	}
	return cols
}

// Close closes the rows iterator.
func (r *Rows) Close() error {
	r.rows.Close()
	return nil
}

// Next is called to populate the next row of data.
func (r *Rows) Next(dest []driver.Value) error {
	if !r.rows.Next() {
		if err := r.rows.Err(); err != nil {
			return err
		}
		return io.EOF
	}
	values, err := r.rows.Values()
	if err != nil {
		return err
	}
	for i, val := range values {
		switch v := val.(type) {
		case int32:
			dest[i] = int64(v)
		case int16:
			dest[i] = int64(v)
		case int8:
			dest[i] = int64(v)
		case int:
			dest[i] = int64(v)
		case uint32:
			dest[i] = int64(v)
		case uint16:
			dest[i] = int64(v)
		case uint8:
			dest[i] = int64(v)
		case uint:
			dest[i] = int64(v)
		case float32:
			dest[i] = float64(v)
		default:
			dest[i] = val
		}
	}
	return nil
}