smart-org/apps/api/app/infrastructure/repositories/credit.py

"""The prepaid credit ledger: every balance change is one atomic update and one history row.

``credit_wallets.balance`` is the number the engine reads; ``credit_transactions``
is the history that explains it. Both change in the caller's transaction, so a
rolled-back request leaves neither behind.
"""

from __future__ import annotations

from datetime import datetime

from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from app.domain.base import utcnow
from app.domain.billing.models import CreditKind, CreditTransaction
from app.infrastructure.db.models import CreditWalletRow
from app.infrastructure.db.models_billing import CreditTransactionRow
from app.infrastructure.repositories.billing import CreditWalletRepository

#: Longest history page one call may return; the UI never needs more at once.
_MAX_HISTORY = 200


def _transaction(row: CreditTransactionRow) -> CreditTransaction:
    return CreditTransaction(
        id=row.id,
        organization_id=row.organization_id,
        kind=CreditKind(row.kind),
        amount=float(row.amount),
        balance_after=float(row.balance_after),
        reason=row.reason,
        actor_user_id=row.actor_user_id,
        usage_event_id=row.usage_event_id,
        created_at=row.created_at,
    )


class CreditLedgerRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def apply(
        self,
        *,
        organization_id: str,
        amount: float,
        kind: CreditKind,
        reason: str = "",
        actor_user_id: str | None = None,
        usage_event_id: str | None = None,
    ) -> CreditTransaction:
        """Move the wallet by a signed ``amount`` and record the move.

        The balance changes in one SQL ``UPDATE`` (``balance = balance + amount``),
        never as a Python read-modify-write, so two writers cannot lose each
        other's change. The history row records the balance that update left.
        """
        await CreditWalletRepository(self.session).ensure(organization_id)
        await self.session.execute(
            update(CreditWalletRow)
            .where(CreditWalletRow.organization_id == organization_id)
            .values(balance=CreditWalletRow.balance + amount, updated_at=utcnow())
            .execution_options(synchronize_session="fetch")
        )
        balance_after = (
            await self.session.execute(
                select(CreditWalletRow.balance).where(
                    CreditWalletRow.organization_id == organization_id
                )
            )
        ).scalar_one()
        row = CreditTransactionRow(
            organization_id=organization_id,
            kind=str(kind),
            amount=amount,
            balance_after=float(balance_after),
            reason=reason[:500],
            actor_user_id=actor_user_id,
            usage_event_id=usage_event_id,
        )
        self.session.add(row)
        await self.session.flush()
        return _transaction(row)

    async def debit_for_usage(
        self, *, organization_id: str, usage_event_id: str, amount: float
    ) -> CreditTransaction | None:
        """Debit the wallet for one usage row, at most once.

        Returns ``None`` when the usage costs nothing or was already debited. The
        debit runs in a savepoint: when two writers race on the same usage row, the
        loser's insert is rolled back and it returns ``None`` instead of poisoning
        the caller's transaction. Any other integrity failure is re-raised.
        """
        if amount <= 0 or await self._debited(organization_id, usage_event_id):
            return None
        try:
            async with self.session.begin_nested():
                return await self.apply(
                    organization_id=organization_id,
                    amount=-amount,
                    kind=CreditKind.USAGE,
                    usage_event_id=usage_event_id,
                )
        except IntegrityError:
            if await self._debited(organization_id, usage_event_id):
                return None
            raise

    async def history(
        self,
        organization_id: str,
        *,
        kinds: tuple[CreditKind, ...] | None = None,
        limit: int = 50,
    ) -> list[CreditTransaction]:
        """One organization's history, newest first, optionally one kind of entry."""
        stmt = select(CreditTransactionRow).where(
            CreditTransactionRow.organization_id == organization_id
        )
        if kinds:
            stmt = stmt.where(CreditTransactionRow.kind.in_([str(k) for k in kinds]))
        stmt = stmt.order_by(CreditTransactionRow.created_at.desc()).limit(
            min(limit, _MAX_HISTORY)
        )
        return [_transaction(r) for r in (await self.session.execute(stmt)).scalars()]

    async def last_credit_at_by_org(self) -> dict[str, datetime]:
        """When each organization last received credit (usage debits do not count).

        Cross-organization on purpose: only the platform owner's views may use it.
        """
        stmt = (
            select(CreditTransactionRow.organization_id, func.max(CreditTransactionRow.created_at))
            .where(CreditTransactionRow.kind != str(CreditKind.USAGE))
            .group_by(CreditTransactionRow.organization_id)
        )
        result = await self.session.execute(stmt)
        return dict(result.tuples().all())

    async def _debited(self, organization_id: str, usage_event_id: str) -> bool:
        stmt = select(CreditTransactionRow.id).where(
            CreditTransactionRow.organization_id == organization_id,
            CreditTransactionRow.usage_event_id == usage_event_id,
        )
        return (await self.session.execute(stmt)).first() is not None