DDI/server/internal/store/store.go

// SPDX-License-Identifier: LicenseRef-DDI-Proprietary

// Package store is the PostgreSQL persistence layer (ADR 0006): connection
// pool, embedded goose migrations and repositories with parameterised SQL.
package store

import (
	"context"
	"embed"
	"errors"
	"fmt"
	"io/fs"
	"log/slog"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/jackc/pgx/v5/stdlib"
	"github.com/pressly/goose/v3"
)

//go:embed migrations/*.sql
var migrationsFS embed.FS

// Store wraps the connection pool.
type Store struct {
	pool *pgxpool.Pool
}

// Open connects to PostgreSQL and verifies the connection.
func Open(ctx context.Context, url string, maxConns int32) (*Store, error) {
	if url == "" {
		return nil, errors.New("no database URL configured")
	}
	cfg, err := pgxpool.ParseConfig(url)
	if err != nil {
		// The parse error may echo the URL including the password.
		return nil, errors.New("invalid database URL")
	}
	cfg.MaxConns = maxConns
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("connect to database: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("connect to database: %w", err)
	}
	return &Store{pool: pool}, nil
}

// New wraps an existing pool (used by tests).
func New(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }

// Close releases all connections.
func (s *Store) Close() { s.pool.Close() }

// Ping checks database connectivity (readiness check).
func (s *Store) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }

// Pool exposes the pool for packages that need direct access (tests).
func (s *Store) Pool() *pgxpool.Pool { return s.pool }

// newProvider creates a goose provider over the embedded migrations.
func (s *Store) newProvider() (*goose.Provider, func() error, error) {
	sub, err := fs.Sub(migrationsFS, "migrations")
	if err != nil {
		return nil, nil, err
	}
	db := stdlib.OpenDBFromPool(s.pool)
	p, err := goose.NewProvider(goose.DialectPostgres, db, sub)
	if err != nil {
		_ = db.Close()
		return nil, nil, err
	}
	return p, db.Close, nil
}

// Migrate applies all pending migrations. It refuses to run against a schema
// that is newer than this binary (downgrade protection).
func (s *Store) Migrate(ctx context.Context, log *slog.Logger) error {
	p, closeDB, err := s.newProvider()
	if err != nil {
		return err
	}
	defer closeDB() //nolint:errcheck // closing the wrapper does not close the pool

	current, err := s.schemaVersion(ctx, p)
	if err != nil {
		return fmt.Errorf("read schema version: %w", err)
	}
	sources := p.ListSources()
	latest := int64(0)
	if len(sources) > 0 {
		latest = sources[len(sources)-1].Version
	}
	if current > latest {
		return fmt.Errorf("database schema version %d is newer than this binary supports (%d); upgrade the server", current, latest)
	}
	results, err := p.Up(ctx)
	if err != nil {
		return fmt.Errorf("apply migrations: %w", err)
	}
	for _, r := range results {
		log.Info("applied migration", slog.Int64("version", r.Source.Version), slog.Duration("duration", r.Duration))
	}
	return nil
}

// SchemaVersion returns the applied schema version.
func (s *Store) SchemaVersion(ctx context.Context) (int64, error) {
	p, closeDB, err := s.newProvider()
	if err != nil {
		return 0, err
	}
	defer closeDB() //nolint:errcheck // closing the wrapper does not close the pool
	return s.schemaVersion(ctx, p)
}

// schemaVersion returns 0 for a database that has never been migrated.
func (s *Store) schemaVersion(ctx context.Context, p *goose.Provider) (int64, error) {
	var exists bool
	if err := s.pool.QueryRow(ctx, `SELECT to_regclass('goose_db_version') IS NOT NULL`).Scan(&exists); err != nil {
		return 0, err
	}
	if !exists {
		return 0, nil
	}
	return p.GetDBVersion(ctx)
}

// querier is implemented by *pgxpool.Pool and pgx.Tx.
type querier 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
	SendBatch(ctx context.Context, b *pgx.Batch) pgx.BatchResults
}

// execBatch sends b and returns the first error.
func (t *Tx) execBatch(ctx context.Context, b *pgx.Batch) error {
	if b.Len() == 0 {
		return nil
	}
	return mapErr(t.q.SendBatch(ctx, b).Close())
}

// Tx is a repository bound to a transaction.
type Tx struct {
	q querier
	// hooks run after a successful commit (nil without a transaction).
	hooks *txHooks
}

type txHooks struct {
	fns  []func()
	keys map[string]bool
}

// Q returns repositories bound to the pool (no transaction).
func (s *Store) Q() *Tx { return &Tx{q: s.pool} }

// InTx runs fn in a transaction that is committed if fn returns nil. The
// functions registered with AfterCommit run after the commit.
func (s *Store) InTx(ctx context.Context, fn func(tx *Tx) error) error {
	h := &txHooks{keys: map[string]bool{}}
	if err := pgx.BeginFunc(ctx, s.pool, func(t pgx.Tx) error { return fn(&Tx{q: t, hooks: h}) }); err != nil {
		return err
	}
	for _, f := range h.fns {
		f()
	}
	return nil
}

// AfterCommit registers fn to run once the transaction has been committed,
// so that notifications never announce changes that are rolled back. A
// non-empty key registers fn only once per transaction (savepoints
// included). Without a transaction fn runs immediately.
func (t *Tx) AfterCommit(key string, fn func()) {
	if t.hooks == nil {
		fn()
		return
	}
	if key != "" {
		if t.hooks.keys[key] {
			return
		}
		t.hooks.keys[key] = true
	}
	t.hooks.fns = append(t.hooks.fns, fn)
}

// Nested runs fn inside a savepoint when t is a transaction. If fn fails,
// only its changes are rolled back and the outer transaction stays usable;
// PostgreSQL would otherwise abort the whole transaction after any error.
func (t *Tx) Nested(ctx context.Context, fn func(tx *Tx) error) error {
	outer, ok := t.q.(pgx.Tx)
	if !ok {
		return fn(t)
	}
	return pgx.BeginFunc(ctx, outer, func(sp pgx.Tx) error { return fn(&Tx{q: sp, hooks: t.hooks}) })
}