DatabaseProgramming: 03 Prevent overpaid invoice trigger.sql

DROP TRIGGER IF EXISTS trg_prevent_overpaid_invoice ON public.payments;
CREATE OR REPLACE FUNCTION public.trg_fn_prevent_overpaid_invoice()
RETURNS trigger
LANGUAGE plpgsql
AS $$
DECLARE
    v_invoice_total numeric;
    v_invoice_status text;
    v_total_paid numeric;
BEGIN
    IF NEW.invoice_id IS NULL OR NEW.status <> 'completed' THEN
        RETURN NEW;
    END IF;

    SELECT total_amount, status
    INTO v_invoice_total, v_invoice_status
    FROM public.invoices
    WHERE invoice_id = NEW.invoice_id
      AND account_id = NEW.account_id
    FOR UPDATE;

    IF NOT FOUND THEN
        RAISE EXCEPTION 'Invoice % does not belong to account %.', NEW.invoice_id, NEW.account_id;
    END IF;
    IF v_invoice_status = 'cancelled' THEN
        RAISE EXCEPTION 'Payment rejected: invoice % is cancelled.', NEW.invoice_id;
    END IF;

    SELECT COALESCE(SUM(amount), 0)
    INTO v_total_paid
    FROM public.payments
    WHERE invoice_id = NEW.invoice_id
      AND status = 'completed'
      AND payment_id IS DISTINCT FROM NEW.payment_id;

    IF v_total_paid >= v_invoice_total - 0.01 THEN
        RAISE EXCEPTION 'Payment rejected: invoice % is already fully paid.', NEW.invoice_id;
    END IF;
    IF v_total_paid + NEW.amount > v_invoice_total + 0.01 THEN
        RAISE EXCEPTION 'Payment rejected: invoice % has remaining balance %, attempted payment is %.',
            NEW.invoice_id, v_invoice_total - v_total_paid, NEW.amount;
    END IF;

    RETURN NEW;
END;
$$;

CREATE TRIGGER trg_prevent_overpaid_invoice
BEFORE INSERT OR UPDATE ON public.payments
FOR EACH ROW
EXECUTE FUNCTION public.trg_fn_prevent_overpaid_invoice();
COMMENT ON FUNCTION public.trg_fn_prevent_overpaid_invoice() IS
'Locks the invoice row and prevents completed payments from exceeding its total, including concurrent attempts.';