refactor(budgets): move the reservation status names beside the column they name

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
 

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))

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",

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``,

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()

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)