X-Financial/server/alembic/versions/20260716_0019_commercial_runtime_reservations.py

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