diff --git a/docs/domains.md b/docs/domains.md index a2a338e7f4..8212675678 100644 --- a/docs/domains.md +++ b/docs/domains.md @@ -142,10 +142,9 @@ Ceilings, reservations, reset periods and per-member policies. - Exceptions: `budget_exceptions.py` - Models: `budgets.py` -`models/budgets.py` holds the scope and reset-alignment vocabularies, with the -columns they name, and the schemas and services both import them from there. -The reservation statuses still sit in `services/budgets/_ledger.py` and move -the same way. +`models/budgets.py` holds the scope, reset-alignment and reservation-status +vocabularies, with the columns they name, and the schemas and services import +them from there. ### pricing diff --git a/src/gateway/models/budgets.py b/src/gateway/models/budgets.py index 4ebdebce9b..bd383b236a 100644 --- a/src/gateway/models/budgets.py +++ b/src/gateway/models/budgets.py @@ -31,7 +31,7 @@ # leaves the sum of three of them ~9000x inside the type. MAX_COUNT_LIMIT = 1_000_000_000_000_000 -# An enum changes the published OpenAPI schema, so both vocabularies stay `Literal`. +# An enum changes the published OpenAPI schema, so the two published vocabularies stay `Literal`. ResetAlignment = Literal["calendar_day", "calendar_week", "calendar_month"] RESET_ALIGNMENTS: tuple[ResetAlignment, ...] = get_args(ResetAlignment) ALIGN_DAY: ResetAlignment = "calendar_day" @@ -46,6 +46,13 @@ SCOPE_ORG_MEMBER: ScopeType = "org_member" SCOPE_API_TOKEN: ScopeType = "api_token" +ReservationStatus = Literal["active", "settled", "released", "expired"] +RESERVATION_STATUSES: tuple[ReservationStatus, ...] = get_args(ReservationStatus) +RESERVATION_ACTIVE: ReservationStatus = "active" +RESERVATION_SETTLED: ReservationStatus = "settled" +RESERVATION_RELEASED: ReservationStatus = "released" +RESERVATION_EXPIRED: ReservationStatus = "expired" + class Budget(Base): """Budget model for spending limits.""" @@ -330,8 +337,10 @@ class BudgetReservation(Base): # takes the hold, and the release has to match what the reserve did. user_reserved: Mapped[bool] = mapped_column(default=False, server_default=false()) # The status is a plain string and not a database enum, so a new state needs no enum migration. - # Its values are the ``RESERVATION_*`` constants in ``services/budgets/_ledger.py``. - status: Mapped[str] = mapped_column(default="active", server_default="active", nullable=False) + # Its values are ``RESERVATION_STATUSES``. + status: Mapped[str] = mapped_column( + default=RESERVATION_ACTIVE, server_default=RESERVATION_ACTIVE, nullable=False + ) # After this instant a still-active row is treated as leaked and reclaimed. expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) diff --git a/src/gateway/services/budgets/_ledger.py b/src/gateway/services/budgets/_ledger.py index 97ed50bfc7..4018b3b847 100644 --- a/src/gateway/services/budgets/_ledger.py +++ b/src/gateway/services/budgets/_ledger.py @@ -22,7 +22,15 @@ from gateway.core.database import create_session from gateway.log_config import logger -from gateway.models.budgets import BudgetReservation, BudgetReservationScope +from gateway.models.budgets import ( + RESERVATION_ACTIVE, + RESERVATION_EXPIRED, + RESERVATION_RELEASED, + RESERVATION_SETTLED, + BudgetReservation, + BudgetReservationScope, + ReservationStatus, +) from gateway.models.users import User from gateway.services.budgets._scoped_enforcement import release as release_scoped @@ -35,13 +43,6 @@ ZERO = Decimal(0) -# The lifecycle, as stored. Plain strings rather than a database enum so a new -# state needs no enum migration (the same reasoning as ``scoped_budgets.scope_type``). -RESERVATION_ACTIVE = "active" -RESERVATION_SETTLED = "settled" # Actual recorded, hold released -RESERVATION_RELEASED = "released" # Hold returned with no spend recorded -RESERVATION_EXPIRED = "expired" # Reclaimed by the TTL sweep after leaking - # The three a row can rest in. Written out as a set the retention query can ask # for by equality: ``status != ACTIVE`` reads the same but is an inequality on # the leading column of ``ix_budget_reservations_status_expires_at``, which the @@ -191,7 +192,7 @@ async def grow( return True -async def try_terminate(db: AsyncSession, reservation_id: str | None, status: str) -> bool: +async def try_terminate(db: AsyncSession, reservation_id: str | None, status: ReservationStatus) -> bool: """Claim the ACTIVE -> terminal transition, reporting whether this caller won. This is the whole point of the ledger. The ``WHERE status = ACTIVE`` guard is @@ -487,10 +488,6 @@ async def run_reservation_sweeper(interval: float, *, batch_size: int, retention __all__ = [ - "RESERVATION_ACTIVE", - "RESERVATION_EXPIRED", - "RESERVATION_RELEASED", - "RESERVATION_SETTLED", "grow", "reclaim_expired_for_user", "prune_terminal", diff --git a/src/gateway/services/budgets/_reservations.py b/src/gateway/services/budgets/_reservations.py index e3b620484b..26e98877e4 100644 --- a/src/gateway/services/budgets/_reservations.py +++ b/src/gateway/services/budgets/_reservations.py @@ -17,7 +17,13 @@ from gateway.core.metered_pricing import estimate_metered_cost from gateway.log_config import logger from gateway.metrics import REGISTRY, Counter -from gateway.models.budgets import MAX_COUNT_LIMIT, Budget, BudgetResetLog +from gateway.models.budgets import ( + MAX_COUNT_LIMIT, + RESERVATION_RELEASED, + RESERVATION_SETTLED, + Budget, + BudgetResetLog, +) from gateway.models.money import to_usd from gateway.models.pricing import ModelPricing from gateway.models.users import User @@ -786,7 +792,7 @@ async def reconcile_reservation( # expression clamps at zero that would pass silently as an under-count of # live holds rather than fail. reclaimed_early = False - if not await ledger.try_terminate(db, handle.reservation_id, ledger.RESERVATION_SETTLED): + if not await ledger.try_terminate(db, handle.reservation_id, RESERVATION_SETTLED): # Losing that claim has two causes and they settle differently. Another # settlement site for this request already ran, and there is nothing left # to do; or the TTL sweep reclaimed the hold while the request was still @@ -899,7 +905,7 @@ async def refund_reservation(db: AsyncSession, handle: ReservationHandle) -> Non reachable from roughly seven sites, and only control flow (a ``raise`` after each) has kept two of them from firing for one request. """ - if not await ledger.try_terminate(db, handle.reservation_id, ledger.RESERVATION_RELEASED): + if not await ledger.try_terminate(db, handle.reservation_id, RESERVATION_RELEASED): # Commit rather than roll back the empty transaction: the guarded UPDATE # matched nothing, so there is nothing to undo, and ``rollback()`` expires # every ORM instance in the session regardless of ``expire_on_commit``, diff --git a/tests/integration/test_budget_reservation_ledger.py b/tests/integration/test_budget_reservation_ledger.py index 9c401c3893..bbacf0c012 100644 --- a/tests/integration/test_budget_reservation_ledger.py +++ b/tests/integration/test_budget_reservation_ledger.py @@ -23,7 +23,16 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from gateway.models.budgets import Budget, BudgetReservation, BudgetReservationScope, ScopedBudget +from gateway.models.budgets import ( + RESERVATION_ACTIVE, + RESERVATION_EXPIRED, + RESERVATION_RELEASED, + RESERVATION_SETTLED, + Budget, + BudgetReservation, + BudgetReservationScope, + ScopedBudget, +) from gateway.models.users import User from gateway.services.budgets import _ledger as ledger from gateway.services.budgets import ( @@ -109,7 +118,7 @@ async def test_a_hold_becomes_a_row_on_both_mechanisms(async_db: AsyncSession, t rows = await _rows(async_db, tenancy.user_id) assert len(rows) == 1 assert rows[0].id == handle.reservation_id - assert rows[0].status == ledger.RESERVATION_ACTIVE + assert rows[0].status == RESERVATION_ACTIVE assert rows[0].user_reserved is True assert rows[0].estimate == Decimal("2.000000") assert rows[0].expires_at > datetime.now(UTC) @@ -145,7 +154,7 @@ async def test_reconcile_is_idempotent_by_reservation_identity(async_db: AsyncSe assert reserved == pytest.approx(0.0) rows = await _rows(async_db, tenancy.user_id) - assert rows[0].status == ledger.RESERVATION_SETTLED + assert rows[0].status == RESERVATION_SETTLED @pytest.mark.asyncio @@ -171,7 +180,7 @@ async def test_refund_is_idempotent_and_records_no_spend(async_db: AsyncSession, assert current == pytest.approx(0.0) assert reserved == pytest.approx(0.0) rows = await _rows(async_db, tenancy.user_id) - assert rows[0].status == ledger.RESERVATION_RELEASED + assert rows[0].status == RESERVATION_RELEASED @pytest.mark.asyncio @@ -217,9 +226,9 @@ async def test_a_second_hold_survives_the_first_being_reclaimed(async_db: AsyncS _, reserved = await _counters(async_db, cap.id) assert reserved == pytest.approx(2.0) - assert await _status(async_db, leaked.reservation_id) == ledger.RESERVATION_EXPIRED + assert await _status(async_db, leaked.reservation_id) == RESERVATION_EXPIRED assert live.reservation_id is not None - assert await _status(async_db, live.reservation_id) == ledger.RESERVATION_ACTIVE + assert await _status(async_db, live.reservation_id) == RESERVATION_ACTIVE @pytest.mark.asyncio @@ -248,7 +257,7 @@ async def test_settling_a_reclaimed_hold_does_not_release_it_twice(async_db: Asy user = await _user(async_db, tenancy.user_id) assert user.spend == Decimal("1.000000") assert user.reserved == Decimal("0.000000") - assert await _status(async_db, handle.reservation_id) == ledger.RESERVATION_SETTLED + assert await _status(async_db, handle.reservation_id) == RESERVATION_SETTLED # And a second late settlement is still a no-op. await reconcile_reservation(async_db, handle, 1.0) @@ -344,7 +353,7 @@ async def test_a_top_up_after_a_reclaim_returns_the_delta(async_db: AsyncSession assert (await _user(async_db, tenancy.user_id)).reserved == Decimal("0.000000") _, reserved = await _counters(async_db, cap.id) assert reserved == pytest.approx(0.0) - assert await _status(async_db, handle.reservation_id) == ledger.RESERVATION_EXPIRED + assert await _status(async_db, handle.reservation_id) == RESERVATION_EXPIRED async_db.expire_all() assert (await async_db.get_one(BudgetReservation, handle.reservation_id)).estimate == Decimal("1.000000") lines = await _lines(async_db, handle.reservation_id) @@ -471,7 +480,7 @@ async def failing_execute(self: AsyncSession, statement: Any, *args: Any, **kwar await async_db.rollback() # The row is still claimable, so the sweep can still return the hold. - assert await _status(async_db, handle.reservation_id) == ledger.RESERVATION_ACTIVE + assert await _status(async_db, handle.reservation_id) == RESERVATION_ACTIVE row = await async_db.get_one(BudgetReservation, handle.reservation_id) row.expires_at = datetime.now(UTC) - timedelta(minutes=1) await async_db.commit() diff --git a/tests/unit/test_budget_vocabularies.py b/tests/unit/test_budget_vocabularies.py index e627cb3e1c..7120c36b0e 100644 --- a/tests/unit/test_budget_vocabularies.py +++ b/tests/unit/test_budget_vocabularies.py @@ -1,4 +1,4 @@ -"""Guards on the two budget vocabularies. +"""Guards on the budget vocabularies. The named constants cover each ``Literal``. The deployment route's scope table has a row for every scope type. @@ -13,6 +13,11 @@ ALIGN_DAY, ALIGN_MONTH, ALIGN_WEEK, + RESERVATION_ACTIVE, + RESERVATION_EXPIRED, + RESERVATION_RELEASED, + RESERVATION_SETTLED, + RESERVATION_STATUSES, RESET_ALIGNMENTS, SCOPE_API_TOKEN, SCOPE_ORG_MEMBER, @@ -33,6 +38,11 @@ def test_the_scope_constants_cover_the_literal() -> None: assert named == set(SCOPE_TYPES) +def test_the_reservation_status_constants_cover_the_literal() -> None: + named = {RESERVATION_ACTIVE, RESERVATION_SETTLED, RESERVATION_RELEASED, RESERVATION_EXPIRED} + assert named == set(RESERVATION_STATUSES) + + def test_the_deployment_route_resolves_every_scope_type() -> None: assert set(_SCOPE_SUBJECTS) == set(SCOPE_TYPES)