"""Repositories for Billing: plans, subscriptions, wallets, invoices, payments. The usage ledger lives in ``repositories/audit.py`` because it is shared with the observability layer; billing reads it but never writes to it directly. """ from __future__ import annotations from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.domain.billing.models import ( CreditWallet, Invoice, InvoiceStatus, Payment, PaymentStatus, Plan, Subscription, SubscriptionStatus, ) from app.infrastructure.db.models import ( CreditWalletRow, InvoiceRow, PaymentRow, PlanRow, SubscriptionRow, ) def _plan(row: PlanRow) -> Plan: return Plan( id=row.id, name=row.name, monthly_price=row.monthly_price, included_credit=row.included_credit, features_json=dict(row.features_json or {}), is_active=row.is_active, created_at=row.created_at, ) def _subscription(row: SubscriptionRow) -> Subscription: return Subscription( id=row.id, organization_id=row.organization_id, plan_id=row.plan_id, status=SubscriptionStatus(row.status), external_subscription_id=row.external_subscription_id, current_period_end=row.current_period_end, created_at=row.created_at, ) def _wallet(row: CreditWalletRow) -> CreditWallet: return CreditWallet( id=row.id, organization_id=row.organization_id, balance=row.balance, currency=row.currency, low_balance_threshold=row.low_balance_threshold, updated_at=row.updated_at, ) def _invoice(row: InvoiceRow) -> Invoice: return Invoice( id=row.id, organization_id=row.organization_id, amount=row.amount, currency=row.currency, status=InvoiceStatus(row.status), external_invoice_id=row.external_invoice_id, period_start=row.period_start, period_end=row.period_end, created_at=row.created_at, ) def _payment(row: PaymentRow) -> Payment: return Payment( id=row.id, organization_id=row.organization_id, amount=row.amount, status=PaymentStatus(row.status), provider=row.provider, external_payment_id=row.external_payment_id, invoice_id=row.invoice_id, created_at=row.created_at, ) class PlanRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def list(self) -> list[Plan]: stmt = select(PlanRow).order_by(PlanRow.monthly_price) return [_plan(r) for r in (await self.session.execute(stmt)).scalars()] async def get(self, plan_id: str) -> Plan | None: row = await self.session.get(PlanRow, plan_id) return _plan(row) if row else None class SubscriptionRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def get_for_org(self, organization_id: str) -> Subscription | None: stmt = select(SubscriptionRow).where( SubscriptionRow.organization_id == organization_id ) row = (await self.session.execute(stmt)).scalars().first() return _subscription(row) if row else None async def upsert(self, subscription: Subscription) -> Subscription: existing = await self.get_for_org(subscription.organization_id) if existing: row = await self.session.get(SubscriptionRow, existing.id) assert row is not None row.plan_id = subscription.plan_id row.status = str(subscription.status) row.external_subscription_id = subscription.external_subscription_id row.current_period_end = subscription.current_period_end await self.session.flush() return _subscription(row) row = SubscriptionRow( id=subscription.id, organization_id=subscription.organization_id, plan_id=subscription.plan_id, status=str(subscription.status), external_subscription_id=subscription.external_subscription_id, current_period_end=subscription.current_period_end, created_at=subscription.created_at, ) self.session.add(row) await self.session.flush() return _subscription(row) class CreditWalletRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def get_for_org(self, organization_id: str) -> CreditWallet | None: stmt = select(CreditWalletRow).where( CreditWalletRow.organization_id == organization_id ) row = (await self.session.execute(stmt)).scalar_one_or_none() return _wallet(row) if row else None async def ensure(self, organization_id: str) -> CreditWallet: """Return the org's wallet, creating a zero-balance one if absent.""" existing = await self.get_for_org(organization_id) if existing: return existing row = CreditWalletRow(organization_id=organization_id, balance=0.0, currency="USD") self.session.add(row) await self.session.flush() return _wallet(row) async def adjust(self, *, organization_id: str, delta: float) -> CreditWallet: """Add (or subtract) credit. Callers must validate the delta's sign.""" wallet = await self.ensure(organization_id) row = await self.session.get(CreditWalletRow, wallet.id) assert row is not None row.balance = float(row.balance) + float(delta) await self.session.flush() return _wallet(row) class InvoiceRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def list_for_org(self, organization_id: str) -> list[Invoice]: stmt = ( select(InvoiceRow) .where(InvoiceRow.organization_id == organization_id) .order_by(InvoiceRow.created_at.desc()) ) return [_invoice(r) for r in (await self.session.execute(stmt)).scalars()] async def create(self, invoice: Invoice) -> Invoice: row = InvoiceRow( id=invoice.id, organization_id=invoice.organization_id, amount=invoice.amount, currency=invoice.currency, status=str(invoice.status), external_invoice_id=invoice.external_invoice_id, period_start=invoice.period_start, period_end=invoice.period_end, created_at=invoice.created_at, ) self.session.add(row) await self.session.flush() return _invoice(row) class PaymentRepository: def __init__(self, session: AsyncSession) -> None: self.session = session async def list_for_org(self, organization_id: str) -> list[Payment]: stmt = ( select(PaymentRow) .where(PaymentRow.organization_id == organization_id) .order_by(PaymentRow.created_at.desc()) ) return [_payment(r) for r in (await self.session.execute(stmt)).scalars()] async def create(self, payment: Payment) -> Payment: row = PaymentRow( id=payment.id, organization_id=payment.organization_id, amount=payment.amount, status=str(payment.status), provider=payment.provider, external_payment_id=payment.external_payment_id, invoice_id=payment.invoice_id, created_at=payment.created_at, ) self.session.add(row) await self.session.flush() return _payment(row) async def set_status( self, payment_id: str, status: PaymentStatus, *, external_id: str | None = None ) -> Payment | None: row = await self.session.get(PaymentRow, payment_id) if row is None: return None row.status = str(status) if external_id is not None: row.external_payment_id = external_id await self.session.flush() return _payment(row)