adapter.go

package pgxadapter

import (
	"context"
	"fmt"
	"runtime"

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

// PgxInterface defines the interface for pgxpool.Pool, pgx.Conn, and pgx.Tx
type PgxInterface interface {
	Exec(ctx context.Context, sql string, arguments ...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
	Begin(ctx context.Context) (pgx.Tx, error)
	CopyFrom(ctx context.Context, tableName pgx.Identifier, columnNames []string, rowSrc pgx.CopyFromSource) (int64, error)
}

type Adapter struct {
	pool      PgxInterface
	tableName string
}

// NewAdapter creates a new adapter with the given connection string.
// It creates the table if it does not exist.
func NewAdapter(connString string, opts ...Option) (*Adapter, error) {
	config, err := pgxpool.ParseConfig(connString)
	if err != nil {
		return nil, err
	}

	pool, err := pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		return nil, err
	}

	return NewAdapterByDB(pool, opts...)
}

// NewAdapterByDB creates a new adapter with the given pgxpool.Pool.
// It creates the table if it does not exist.
func NewAdapterByDB(pool *pgxpool.Pool, opts ...Option) (*Adapter, error) {
	a := &Adapter{
		pool:      pool,
		tableName: "casbin_rule",
	}

	a.applyOptions(opts...)

	if err := a.createTable(); err != nil {
		return nil, err
	}

	return a, nil
}

// NewAdapterByConn creates a new adapter with the given pgx.Conn.
// It creates the table if it does not exist.
func NewAdapterByConn(conn *pgx.Conn, opts ...Option) (*Adapter, error) {
	a := &Adapter{
		pool:      conn,
		tableName: "casbin_rule",
	}

	a.applyOptions(opts...)

	if err := a.createTable(); err != nil {
		return nil, err
	}

	return a, nil
}

// NewAdapterByDBWithCustomTable creates a new adapter with the given pgxpool.Pool and custom table name.
// It creates the table if it does not exist.
//
// Deprecated: Use NewAdapterByDB with WithTableName option instead.
func NewAdapterByDBWithCustomTable(pool *pgxpool.Pool, tableName string) (*Adapter, error) {
	return NewAdapterByDB(pool, WithTableName(tableName))
}

func (a *Adapter) createTable() error {
	_, err := a.pool.Exec(context.Background(), fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s (
			id SERIAL PRIMARY KEY,
			ptype VARCHAR(100),
			v0 VARCHAR(100),
			v1 VARCHAR(100),
			v2 VARCHAR(100),
			v3 VARCHAR(100),
			v4 VARCHAR(100),
			v5 VARCHAR(100)
		)`, a.tableName))
	return err
}

func finalizer(a *Adapter) {
	a.Close()
}

// Close closes the connection pool or connection.
func (a *Adapter) Close() {
	if a.pool != nil {
		if pool, ok := a.pool.(*pgxpool.Pool); ok {
			pool.Close()
		} else if conn, ok := a.pool.(*pgx.Conn); ok {
			conn.Close(context.Background())
		}
	}
}

// SetFinalizer sets the finalizer for the adapter.
// This is optional and mostly for cleanup if the user forgets to close.
func (a *Adapter) SetFinalizer() {
	runtime.SetFinalizer(a, finalizer)
}