textmachine/platform/internal/pgstore/credits.go

package pgstore

import (
	"context"
	"errors"
	"fmt"
	"time"

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

	"textmachine/platform/internal/money"
)

var (
	// ErrInsufficientCredit is a refusal, not a failure: the account has less than the run needs.
	ErrInsufficientCredit = errors.New("pgstore: insufficient credit")
	// ErrNoReservation means the hold this settlement refers to is not open.
	ErrNoReservation = errors.New("pgstore: no open reservation")
	// ErrDuplicateHold is a second hold on an attempt id that already has one. It is an error, not
	// a no-op: a hold that debits nothing reserves nothing while reporting that it did.
	ErrDuplicateHold = errors.New("pgstore: attempt already has a hold")
	// ErrNotOwner is a book that does not belong to the account being charged for it.
	ErrNotOwner = errors.New("pgstore: book belongs to another account")
	// ErrNoAccount separates "this account has nothing" from "this account does not exist" — the
	// difference between a balance of zero and a typo in an admin command.
	ErrNoAccount = errors.New("pgstore: no such account")
)

// Grant credits an account and reports whether this call is what credited it. The free tier is one
// of these and nothing more.
//
// (source, sourceID) is the idempotency key, scoped to the account by the schema. applied is false
// when the key was already spent: the caller must say so rather than print a success it did not
// cause.
func (s *Store) Grant(ctx context.Context, userID string, amount money.MicroUSD, source, sourceID, note string, now time.Time) (applied bool, err error) {
	if amount <= 0 {
		return false, fmt.Errorf("pgstore: grant must be positive, got %d", amount)
	}
	if source == "" || sourceID == "" {
		return false, errors.New("pgstore: grant needs an idempotency key")
	}
	err = s.inTx(ctx, func(tx pgx.Tx) error {
		applied, err = appendLedger(ctx, tx, userID, "grant", amount, source, sourceID, note, now)
		return err
	})
	return applied, err
}

// Adjust corrects a balance. A ledger row is never edited: the correction is another row, which is
// what keeps the sum reproducible. The note is mandatory, in the DDL as well as here.
func (s *Store) Adjust(ctx context.Context, userID string, amount money.MicroUSD, source, sourceID, note string, now time.Time) (applied bool, err error) {
	if amount == 0 || note == "" {
		return false, errors.New("pgstore: an adjustment needs a non-zero amount and a reason")
	}
	if source == "" || sourceID == "" {
		return false, errors.New("pgstore: adjustment needs an idempotency key")
	}
	err = s.inTx(ctx, func(tx pgx.Tx) error {
		applied, err = appendLedger(ctx, tx, userID, "adjustment", amount, source, sourceID, note, now)
		return err
	})
	return applied, err
}

// Balance is what the account may still spend, read from the cache that every ledger write updates
// in its own transaction.
func (s *Store) Balance(ctx context.Context, userID string) (money.MicroUSD, error) {
	var v *int64
	err := s.pool.QueryRow(ctx, `
		select b.balance_micro_usd
		  from users u left join account_balances b on b.user_id = u.id
		 where u.id = $1`, userID).Scan(&v)
	if errors.Is(err, pgx.ErrNoRows) {
		return 0, ErrNoAccount
	}
	if err != nil {
		return 0, fmt.Errorf("pgstore: balance: %w", err)
	}
	if v == nil {
		return 0, nil // an account with no ledger rows has no credit, which is not an error
	}
	return money.MicroUSD(*v), nil
}

// Account is what an operator needs to see about one account's money.
type Account struct {
	Balance  money.MicroUSD
	Reserved money.MicroUSD
	// LedgerSum is recomputed from the rows. It exists to be COMPARED with Balance, and both are
	// read in one snapshot: reading them separately reports drift that a concurrent grant caused
	// between the two queries.
	LedgerSum money.MicroUSD
}

// ReadAccount returns the money view in a single consistent snapshot.
func (s *Store) ReadAccount(ctx context.Context, userID string) (Account, error) {
	var a Account
	const q = `
		select coalesce((select balance_micro_usd from account_balances where user_id = $1), 0),
		       coalesce((select sum(amount_micro_usd) from credit_ledger where user_id = $1), 0),
		       coalesce((select sum(amount_micro_usd) from reservations
		                  where user_id = $1 and state = 'open'), 0)
		  from users where id = $1`
	err := s.pool.QueryRow(ctx, q, userID).Scan(&a.Balance, &a.LedgerSum, &a.Reserved)
	if errors.Is(err, pgx.ErrNoRows) {
		return Account{}, ErrNoAccount
	}
	if err != nil {
		return Account{}, fmt.Errorf("pgstore: read account: %w", err)
	}
	return a, nil
}

// Reservation is an open hold as an operator sees it.
type Reservation struct {
	EngineRunID string
	BookID      string
	Amount      money.MicroUSD
	Ceiling     money.MicroUSD
	OpenedAt    time.Time
}

// OpenReservations lists holds that were taken and never closed. Without this they are money that
// is gone from the balance and invisible to everything that could give it back.
func (s *Store) OpenReservations(ctx context.Context, userID string) ([]Reservation, error) {
	const q = `
		select engine_run_id, book_id, amount_micro_usd, ceiling_micro_usd, opened_at
		  from reservations where user_id = $1 and state = 'open' order by opened_at`
	rows, err := s.pool.Query(ctx, q, userID)
	if err != nil {
		return nil, fmt.Errorf("pgstore: open reservations: %w", err)
	}
	defer rows.Close()
	var out []Reservation
	for rows.Next() {
		var r Reservation
		if err := rows.Scan(&r.EngineRunID, &r.BookID, &r.Amount, &r.Ceiling, &r.OpenedAt); err != nil {
			return nil, fmt.Errorf("pgstore: scan reservation: %w", err)
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// Hold reserves credit before a run is spawned. Together with the per-book ceiling handed to the
// engine it is the enforcement half of the money design: the hold makes the credit unavailable to
// the next run, and the engine stops itself at the ceiling, so an overspend is impossible even
// while the platform is blind. The event stream is freshness only.
//
// The ceiling to hand the engine is the amount held; read it back with OpenReservations.
func (s *Store) Hold(ctx context.Context, userID, bookID, engineRunID string, amount money.MicroUSD, now time.Time) error {
	if amount <= 0 {
		return fmt.Errorf("pgstore: hold must be positive, got %d", amount)
	}
	return s.inTx(ctx, func(tx pgx.Tx) error {
		// account_balances is locked FIRST here and in every other operation. A path that locked
		// the reservation first would invert the order against this one and deadlock — measured,
		// not hypothetical.
		balance, err := lockBalance(ctx, tx, userID)
		if err != nil {
			return err
		}
		if balance < amount {
			return ErrInsufficientCredit
		}
		// The book must belong to the account being charged. The foreign key only proves the book
		// exists, which is not the same question.
		tag, err := tx.Exec(ctx, `
			insert into reservations (engine_run_id, user_id, book_id, amount_micro_usd, ceiling_micro_usd, state, opened_at)
			select $1, $2, $3, $4, $4, 'open', $5 from books where id = $3 and owner_id = $2`,
			engineRunID, userID, bookID, int64(amount), now)
		if err != nil {
			return fmt.Errorf("pgstore: open reservation: %w", err)
		}
		if tag.RowsAffected() == 0 {
			return ErrNotOwner
		}
		applied, err := appendLedger(ctx, tx, userID, "hold", -amount, "run", engineRunID, "", now)
		if err != nil {
			return err
		}
		if !applied {
			// The ledger already holds this attempt id, so nothing was debited. Reporting success
			// would spawn a run against credit that was never reserved.
			return ErrDuplicateHold
		}
		return nil
	})
}

// Settle closes a reservation with what the attempt actually cost: the hold comes back and the real
// cost is charged, in one transaction. Settling twice is refused — the reservation is no longer
// open — which is what makes it safe on a retried path.
//
// A cost above the hold is CAPPED at the hold and the row says so. Spending more than was reserved
// means the engine's ceiling did not hold, and the account is not the place to absorb that.
func (s *Store) Settle(ctx context.Context, engineRunID string, spent money.MicroUSD, now time.Time) error {
	if spent < 0 {
		return fmt.Errorf("pgstore: spend cannot be negative, got %d", spent)
	}
	return s.inTx(ctx, func(tx pgx.Tx) error {
		userID, held, err := closeReservation(ctx, tx, engineRunID, "settled", now)
		if err != nil {
			return err
		}
		note := ""
		if spent > held {
			note = fmt.Sprintf("capped at the hold; the engine reported %s", spent.USD())
			spent = held
		}
		if _, err := appendLedger(ctx, tx, userID, "hold_release", held, "run_release", engineRunID, "", now); err != nil {
			return err
		}
		_, err = appendLedger(ctx, tx, userID, "settlement", -spent, "run_settle", engineRunID, note, now)
		return err
	})
}

// Release gives a reservation back untouched: the run never started, or it cost nothing.
func (s *Store) Release(ctx context.Context, engineRunID string, now time.Time) error {
	return s.inTx(ctx, func(tx pgx.Tx) error {
		userID, held, err := closeReservation(ctx, tx, engineRunID, "released", now)
		if err != nil {
			return err
		}
		_, err = appendLedger(ctx, tx, userID, "hold_release", held, "run_release", engineRunID, "", now)
		return err
	})
}

// lockBalance takes the account's row lock and returns the balance under it. Every money operation
// starts here, so they all take their locks in the same order.
func lockBalance(ctx context.Context, tx pgx.Tx, userID string) (money.MicroUSD, error) {
	var v int64
	err := tx.QueryRow(ctx, `select balance_micro_usd from account_balances where user_id = $1 for update`, userID).Scan(&v)
	if errors.Is(err, pgx.ErrNoRows) {
		return 0, ErrInsufficientCredit // no ledger row yet means no credit
	}
	if err != nil {
		return 0, fmt.Errorf("pgstore: read balance: %w", err)
	}
	return money.MicroUSD(v), nil
}

func closeReservation(ctx context.Context, tx pgx.Tx, engineRunID, state string, now time.Time) (string, money.MicroUSD, error) {
	// Read the owner unlocked, lock the balance, then close under the state guard. The guard is
	// what makes the unlocked read safe: a reservation closed by someone else in between makes the
	// update match nothing.
	var userID string
	err := tx.QueryRow(ctx, `select user_id from reservations where engine_run_id = $1`, engineRunID).Scan(&userID)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", 0, ErrNoReservation
	}
	if err != nil {
		return "", 0, fmt.Errorf("pgstore: find reservation: %w", err)
	}
	if _, err := lockBalance(ctx, tx, userID); err != nil && !errors.Is(err, ErrInsufficientCredit) {
		return "", 0, err
	}
	var amount int64
	err = tx.QueryRow(ctx, `
		update reservations set state = $2, closed_at = $3
		 where engine_run_id = $1 and state = 'open'
		returning amount_micro_usd`, engineRunID, state, now).Scan(&amount)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", 0, ErrNoReservation
	}
	if err != nil {
		return "", 0, fmt.Errorf("pgstore: close reservation: %w", err)
	}
	return userID, money.MicroUSD(amount), nil
}

// appendLedger writes one row and moves the cached balance with it, in the caller's transaction.
// The two are never written apart: a cache that can lag its source is a second answer about money.
// applied is false when the idempotency key was already spent.
func appendLedger(ctx context.Context, tx pgx.Tx, userID, kind string, amount money.MicroUSD, source, sourceID, note string, now time.Time) (bool, error) {
	tag, err := tx.Exec(ctx, `
		insert into credit_ledger (user_id, kind, amount_micro_usd, source, source_id, note, created_at)
		values ($1, $2, $3, $4, $5, $6, $7)
		on conflict (user_id, source, source_id) do nothing`,
		userID, kind, int64(amount), source, sourceID, note, now)
	if err != nil {
		// A typo in an account id is the commonest way an operator gets here, and Balance already
		// answers it with ErrNoAccount. Reporting the same fact as a raw constraint name reads as a
		// broken database (PD-56).
		var pg *pgconn.PgError
		if errors.As(err, &pg) && pg.ConstraintName == "credit_ledger_user_id_fkey" {
			return false, ErrNoAccount
		}
		return false, fmt.Errorf("pgstore: append ledger: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return false, nil
	}
	if _, err := tx.Exec(ctx, `
		insert into account_balances (user_id, balance_micro_usd, updated_at)
		values ($1, $2, $3)
		on conflict (user_id) do update
		   set balance_micro_usd = account_balances.balance_micro_usd + excluded.balance_micro_usd,
		       updated_at = excluded.updated_at`,
		userID, int64(amount), now); err != nil {
		return false, fmt.Errorf("pgstore: update balance: %w", err)
	}
	return true, nil
}

func (s *Store) inTx(ctx context.Context, fn func(pgx.Tx) error) error {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("pgstore: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if err := fn(tx); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("pgstore: commit: %w", err)
	}
	return nil
}

// DeleteBook removes a book and the CLOSED reservations that referenced it. An OPEN one blocks the
// delete (the foreign key is RESTRICT), which is the point: removing a book with money reserved
// against it would leave the hold in the ledger with nothing left to release it.
//
// Closed reservations carry no financial fact the ledger does not already hold — they are
// operational state — so removing them with the book loses nothing.
func (s *Store) DeleteBook(ctx context.Context, bookID string) error {
	return s.inTx(ctx, func(tx pgx.Tx) error {
		if _, err := tx.Exec(ctx,
			`delete from reservations where book_id = $1 and state <> 'open'`, bookID); err != nil {
			return fmt.Errorf("pgstore: clear reservations: %w", err)
		}
		if _, err := tx.Exec(ctx, `delete from books where id = $1`, bookID); err != nil {
			return fmt.Errorf("pgstore: delete book: %w", err)
		}
		return nil
	})
}