787bc3a481
Add tenant-safe value, telemetry, connector, commercial, and production-readiness foundations.
83 lines
2.1 KiB
Python
83 lines
2.1 KiB
Python
"""expand commercial runtime resource quantity bases
|
|
|
|
Revision ID: 20260717_0024
|
|
Revises: 20260716_0023
|
|
Create Date: 2026-07-17 09:30:00
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import sqlalchemy as sa
|
|
|
|
from alembic import op
|
|
|
|
revision: str = "20260717_0024"
|
|
down_revision: str | None = "20260716_0023"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
_CONSTRAINT = "ck_commercial_runtime_reservations_basis"
|
|
_LEGACY_BASES = (
|
|
"quantity_basis IN ('call', 'input_tokens', 'output_tokens', "
|
|
"'total_tokens', 'duration_ms')"
|
|
)
|
|
_RESOURCE_BASES = (
|
|
"quantity_basis IN ('call', 'input_tokens', 'output_tokens', "
|
|
"'total_tokens', 'duration_ms', 'bytes', 'pages', 'objects', 'events')"
|
|
)
|
|
|
|
|
|
def _require_postgresql() -> None:
|
|
dialect_name = op.get_bind().dialect.name
|
|
if dialect_name != "postgresql":
|
|
raise RuntimeError(
|
|
"20260717_0024 only supports PostgreSQL; "
|
|
f"refusing to mutate {dialect_name} without transactional constraint DDL"
|
|
)
|
|
|
|
|
|
def _require_no_resource_reservations_for_downgrade() -> None:
|
|
count = int(
|
|
op.get_bind().scalar(
|
|
sa.text(
|
|
"SELECT COUNT(*) FROM commercial_runtime_reservations "
|
|
"WHERE quantity_basis IN ('bytes', 'pages', 'objects', 'events')"
|
|
)
|
|
)
|
|
or 0
|
|
)
|
|
if count:
|
|
raise RuntimeError(
|
|
"cannot downgrade commercial resource quantity bases: "
|
|
f"resource reservations exist ({count})"
|
|
)
|
|
|
|
|
|
def upgrade() -> None:
|
|
_require_postgresql()
|
|
op.drop_constraint(
|
|
_CONSTRAINT,
|
|
"commercial_runtime_reservations",
|
|
type_="check",
|
|
)
|
|
op.create_check_constraint(
|
|
_CONSTRAINT,
|
|
"commercial_runtime_reservations",
|
|
_RESOURCE_BASES,
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
_require_postgresql()
|
|
_require_no_resource_reservations_for_downgrade()
|
|
op.drop_constraint(
|
|
_CONSTRAINT,
|
|
"commercial_runtime_reservations",
|
|
type_="check",
|
|
)
|
|
op.create_check_constraint(
|
|
_CONSTRAINT,
|
|
"commercial_runtime_reservations",
|
|
_LEGACY_BASES,
|
|
)
|