"""add commercial runtime quota reservations Revision ID: 20260716_0019 Revises: 20260716_0018 Create Date: 2026-07-16 22:10:00 """ from collections.abc import Sequence import sqlalchemy as sa from alembic import op revision: str = "20260716_0019" down_revision: str | None = "20260716_0018" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None def _require_postgresql() -> None: dialect_name = op.get_bind().dialect.name if dialect_name != "postgresql": raise RuntimeError( "20260716_0019 only supports PostgreSQL; " f"refusing to mutate {dialect_name} without transactional constraint DDL" ) def _require_empty_reservations_for_downgrade() -> None: count = int( op.get_bind().scalar( sa.text("SELECT COUNT(*) FROM commercial_runtime_reservations") ) or 0 ) if count: raise RuntimeError( "cannot downgrade commercial runtime reservations: " f"operational quota holds exist ({count})" ) def upgrade() -> None: _require_postgresql() op.create_table( "commercial_runtime_reservations", sa.Column("id", sa.String(length=36), nullable=False), sa.Column("tenant_id", sa.String(length=64), nullable=False), sa.Column("subscription_id", sa.String(length=36), nullable=False), sa.Column("entitlement_id", sa.String(length=36), nullable=False), sa.Column("run_id", sa.String(length=50), nullable=False), sa.Column("tool_call_id", sa.String(length=36), nullable=False), sa.Column("tool_type", sa.String(length=30), nullable=False), sa.Column("tool_name", sa.String(length=100), nullable=False), sa.Column("quantity_basis", sa.String(length=20), nullable=False), sa.Column("reserved_quantity", sa.Numeric(20, 6), nullable=False), sa.Column("actual_quantity", sa.Numeric(20, 6), nullable=True), sa.Column("period_key", sa.String(length=32), nullable=False), sa.Column("status", sa.String(length=32), nullable=False), sa.Column("request_fingerprint", sa.String(length=64), nullable=False), sa.Column("meter_config_json", sa.JSON(), nullable=False), sa.Column("resolution_code", sa.String(length=64), nullable=True), sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False), sa.Column("settled_at", sa.DateTime(timezone=True), nullable=True), sa.Column( "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.Column( "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False, ), sa.CheckConstraint( "status IN ('reserved', 'committed', 'released', 'expired', " "'reconciliation_required', 'committed_reconciliation_required')", name="ck_commercial_runtime_reservations_status", ), sa.CheckConstraint( "quantity_basis IN ('call', 'input_tokens', 'output_tokens', " "'total_tokens', 'duration_ms')", name="ck_commercial_runtime_reservations_basis", ), sa.CheckConstraint( "reserved_quantity > 0 AND " "(actual_quantity IS NULL OR actual_quantity > 0)", name="ck_commercial_runtime_reservations_quantity", ), sa.CheckConstraint( "expires_at > created_at", name="ck_commercial_runtime_reservations_expiry", ), sa.CheckConstraint( "(status = 'reserved' AND actual_quantity IS NULL AND settled_at IS NULL " "AND resolution_code IS NULL) OR " "(status = 'committed' AND actual_quantity IS NOT NULL " "AND actual_quantity <= reserved_quantity AND settled_at IS NOT NULL " "AND resolution_code IS NULL) OR " "(status IN ('released', 'expired') AND actual_quantity IS NULL " "AND settled_at IS NOT NULL AND resolution_code IS NOT NULL) OR " "(status = 'reconciliation_required' AND settled_at IS NULL " "AND resolution_code IS NOT NULL) OR " "(status = 'committed_reconciliation_required' " "AND actual_quantity IS NOT NULL " "AND actual_quantity <= reserved_quantity " "AND settled_at IS NOT NULL AND resolution_code IS NOT NULL)", name="ck_commercial_runtime_reservations_state", ), sa.CheckConstraint( "length(trim(run_id)) > 0 AND length(trim(tool_call_id)) > 0 " "AND length(trim(tool_type)) > 0 AND length(trim(tool_name)) > 0 " "AND length(trim(period_key)) > 0 " "AND length(trim(request_fingerprint)) > 0", name="ck_commercial_runtime_reservations_keys", ), sa.ForeignKeyConstraint( ["tenant_id", "subscription_id"], ["tenant_subscriptions.tenant_id", "tenant_subscriptions.id"], name="fk_commercial_runtime_reservations_tenant_subscription", ondelete="RESTRICT", ), sa.ForeignKeyConstraint( ["tenant_id", "subscription_id", "entitlement_id"], [ "commercial_entitlements.tenant_id", "commercial_entitlements.subscription_id", "commercial_entitlements.id", ], name="fk_commercial_runtime_reservations_tenant_entitlement", ondelete="RESTRICT", ), sa.PrimaryKeyConstraint("id"), sa.UniqueConstraint( "tenant_id", "id", name="uq_commercial_runtime_reservations_tenant_id", ), sa.UniqueConstraint( "tool_call_id", name="uq_commercial_runtime_reservations_tool_call", ), ) op.create_index( "ix_commercial_runtime_reservations_quota", "commercial_runtime_reservations", ["tenant_id", "subscription_id", "entitlement_id", "period_key", "status"], ) op.create_index( "ix_commercial_runtime_reservations_expiry", "commercial_runtime_reservations", ["status", "expires_at"], ) op.create_index( "ix_commercial_runtime_reservations_run", "commercial_runtime_reservations", ["tenant_id", "run_id", "created_at"], ) def downgrade() -> None: _require_postgresql() _require_empty_reservations_for_downgrade() op.drop_index( "ix_commercial_runtime_reservations_run", table_name="commercial_runtime_reservations", ) op.drop_index( "ix_commercial_runtime_reservations_expiry", table_name="commercial_runtime_reservations", ) op.drop_index( "ix_commercial_runtime_reservations_quota", table_name="commercial_runtime_reservations", ) op.drop_table("commercial_runtime_reservations")