video-gen/video-gen-api/alembic/versions/20da1d353914_repair_llm_billing_optimize_idempotency.py

"""repair llm billing optimize idempotency

Revision ID: 20da1d353914
Revises: 6a3ea8d0b4c8
Create Date: 2026-07-24 13:35:48.513710
"""
from __future__ import annotations

import re
from typing import Any, Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.engine import Connection


# revision identifiers, used by Alembic.
revision: str = "20da1d353914"
down_revision: Union[str, None] = "6a3ea8d0b4c8"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


_PROMPT_USAGE_COLUMN = "prompt_usage_snapshot_json"

_INDEX_SPECS: tuple[dict[str, Any], ...] = (
    {
        "name": "uq_credit_records_user_refund_target",
        "table": "credit_records",
        "columns": ("user_id", "refund_for_biz_key"),
        "where_sql": "type = 'refund' AND refund_for_biz_key IS NOT NULL",
        "definition_fragments": (
            "(user_id, refund_for_biz_key)",
            "type",
            "refund",
            "refund_for_biz_key is not null",
        ),
        "duplicate_sql": """
            SELECT
                user_id,
                refund_for_biz_key,
                COUNT(*) AS duplicate_count
            FROM credit_records
            WHERE type = 'refund'
              AND refund_for_biz_key IS NOT NULL
            GROUP BY user_id, refund_for_biz_key
            HAVING COUNT(*) > 1
            ORDER BY duplicate_count DESC, user_id, refund_for_biz_key
            LIMIT 20
        """,
    },
    {
        "name": "uq_genrec_user_idempotency_active",
        "table": "generation_records",
        "columns": ("user_id", "idempotency_key"),
        "where_sql": "idempotency_key IS NOT NULL AND deleted_at IS NULL",
        "definition_fragments": (
            "(user_id, idempotency_key)",
            "idempotency_key is not null",
            "deleted_at is null",
        ),
        "duplicate_sql": """
            SELECT
                user_id,
                idempotency_key,
                COUNT(*) AS duplicate_count
            FROM generation_records
            WHERE idempotency_key IS NOT NULL
              AND deleted_at IS NULL
            GROUP BY user_id, idempotency_key
            HAVING COUNT(*) > 1
            ORDER BY duplicate_count DESC, user_id, idempotency_key
            LIMIT 20
        """,
    },
    {
        "name": "uq_token_usage_user_biz_key",
        "table": "token_usage",
        "columns": ("user_id", "biz_key"),
        "where_sql": "biz_key IS NOT NULL",
        "definition_fragments": (
            "(user_id, biz_key)",
            "biz_key is not null",
        ),
        "duplicate_sql": """
            SELECT
                user_id,
                biz_key,
                COUNT(*) AS duplicate_count
            FROM token_usage
            WHERE user_id IS NOT NULL
              AND biz_key IS NOT NULL
            GROUP BY user_id, biz_key
            HAVING COUNT(*) > 1
            ORDER BY duplicate_count DESC, user_id, biz_key
            LIMIT 20
        """,
    },
)


def _bind() -> Connection:
    bind = op.get_bind()
    if bind.dialect.name != "postgresql":
        raise RuntimeError(
            "Migration 20da1d353914 requires PostgreSQL; "
            f"current dialect is {bind.dialect.name!r}."
        )
    return bind


def _current_schema(bind: Connection) -> str:
    schema = bind.execute(sa.text("SELECT current_schema()")).scalar_one_or_none()
    if not schema:
        raise RuntimeError("Unable to resolve PostgreSQL current_schema().")
    return str(schema)


def _require_table(bind: Connection, schema: str, table_name: str) -> None:
    exists = bind.execute(
        sa.text(
            """
            SELECT EXISTS (
                SELECT 1
                FROM information_schema.tables
                WHERE table_schema = :schema
                  AND table_name = :table_name
            )
            """
        ),
        {"schema": schema, "table_name": table_name},
    ).scalar_one()
    if not bool(exists):
        raise RuntimeError(
            f"Required table {schema}.{table_name} does not exist; "
            "refusing to apply migration on an unexpected schema baseline."
        )


def _column_exists(
    bind: Connection,
    schema: str,
    table_name: str,
    column_name: str,
) -> bool:
    return bool(
        bind.execute(
            sa.text(
                """
                SELECT EXISTS (
                    SELECT 1
                    FROM information_schema.columns
                    WHERE table_schema = :schema
                      AND table_name = :table_name
                      AND column_name = :column_name
                )
                """
            ),
            {
                "schema": schema,
                "table_name": table_name,
                "column_name": column_name,
            },
        ).scalar_one()
    )


def _index_info(
    bind: Connection,
    schema: str,
    index_name: str,
) -> dict[str, Any] | None:
    row = bind.execute(
        sa.text(
            """
            SELECT
                table_rel.relname AS table_name,
                index_meta.indisvalid AS is_valid,
                index_meta.indisunique AS is_unique,
                pg_get_indexdef(index_rel.oid) AS index_definition
            FROM pg_class AS index_rel
            JOIN pg_namespace AS namespace_rel
              ON namespace_rel.oid = index_rel.relnamespace
            JOIN pg_index AS index_meta
              ON index_meta.indexrelid = index_rel.oid
            JOIN pg_class AS table_rel
              ON table_rel.oid = index_meta.indrelid
            WHERE namespace_rel.nspname = :schema
              AND index_rel.relname = :index_name
            """
        ),
        {"schema": schema, "index_name": index_name},
    ).mappings().one_or_none()
    return dict(row) if row is not None else None


def _normalize_index_definition(value: str) -> str:
    normalized = value.lower().replace('"', "")
    normalized = re.sub(r"::[a-z_ ]+(?:\[\])?", "", normalized)
    normalized = re.sub(r"[()]", lambda match: match.group(0), normalized)
    normalized = re.sub(r"\s+", " ", normalized)
    return normalized.strip()


def _assert_existing_index_matches(spec: dict[str, Any], info: dict[str, Any]) -> None:
    if str(info["table_name"]) != str(spec["table"]):
        raise RuntimeError(
            f"Index {spec['name']} already exists on table {info['table_name']}, "
            f"expected table {spec['table']}."
        )
    if not bool(info["is_unique"]):
        raise RuntimeError(
            f"Index {spec['name']} already exists but is not UNIQUE."
        )

    definition = _normalize_index_definition(str(info["index_definition"] or ""))
    missing = [
        fragment
        for fragment in spec["definition_fragments"]
        if fragment not in definition
    ]
    if missing:
        raise RuntimeError(
            f"Index {spec['name']} already exists with an unexpected definition; "
            f"missing expected fragments: {missing}. Actual definition: "
            f"{info['index_definition']}"
        )


def _assert_no_duplicates(bind: Connection, spec: dict[str, Any]) -> None:
    rows = bind.execute(sa.text(spec["duplicate_sql"])).mappings().all()
    if not rows:
        return

    samples = "; ".join(
        ", ".join(f"{key}={value!r}" for key, value in row.items())
        for row in rows
    )
    raise RuntimeError(
        f"Cannot create unique index {spec['name']}: duplicate historical data "
        f"exists. Resolve the conflicts first. Samples: {samples}"
    )


def _quote_identifier(bind: Connection, value: str) -> str:
    return bind.dialect.identifier_preparer.quote(value)


def _create_missing_indexes_concurrently(
    bind: Connection,
    schema: str,
) -> None:
    missing_specs: list[dict[str, Any]] = []

    for spec in _INDEX_SPECS:
        _require_table(bind, schema, str(spec["table"]))
        info = _index_info(bind, schema, str(spec["name"]))
        if info is None:
            _assert_no_duplicates(bind, spec)
            missing_specs.append(spec)
            continue

        if bool(info["is_valid"]):
            _assert_existing_index_matches(spec, info)
            continue

        # A failed CREATE INDEX CONCURRENTLY may leave an invalid index behind.
        # Remove only the invalid index with this migration-owned name, then retry.
        missing_specs.append(spec)

    if not missing_specs:
        return

    context = op.get_context()
    with context.autocommit_block():
        for spec in missing_specs:
            current = _index_info(bind, schema, str(spec["name"]))
            quoted_index = _quote_identifier(bind, str(spec["name"]))
            quoted_schema = _quote_identifier(bind, schema)
            if current is not None:
                if bool(current["is_valid"]):
                    _assert_existing_index_matches(spec, current)
                    continue
                op.execute(
                    sa.text(
                        f"DROP INDEX CONCURRENTLY IF EXISTS "
                        f"{quoted_schema}.{quoted_index}"
                    )
                )

            quoted_table = _quote_identifier(bind, str(spec["table"]))
            quoted_columns = ", ".join(
                _quote_identifier(bind, str(column))
                for column in spec["columns"]
            )
            op.execute(
                sa.text(
                    f"CREATE UNIQUE INDEX CONCURRENTLY {quoted_index} "
                    f"ON {quoted_schema}.{quoted_table} ({quoted_columns}) "
                    f"WHERE {spec['where_sql']}"
                )
            )


def _drop_indexes_concurrently(bind: Connection, schema: str) -> None:
    existing_names = [
        str(spec["name"])
        for spec in _INDEX_SPECS
        if _index_info(bind, schema, str(spec["name"])) is not None
    ]
    if not existing_names:
        return

    quoted_schema = _quote_identifier(bind, schema)
    context = op.get_context()
    with context.autocommit_block():
        for index_name in existing_names:
            quoted_index = _quote_identifier(bind, index_name)
            op.execute(
                sa.text(
                    f"DROP INDEX CONCURRENTLY IF EXISTS "
                    f"{quoted_schema}.{quoted_index}"
                )
            )


def upgrade() -> None:
    bind = _bind()
    schema = _current_schema(bind)

    _require_table(bind, schema, "generation_records")
    if not _column_exists(
        bind,
        schema,
        "generation_records",
        _PROMPT_USAGE_COLUMN,
    ):
        op.add_column(
            "generation_records",
            sa.Column(_PROMPT_USAGE_COLUMN, sa.Text(), nullable=True),
            schema=schema,
        )

    _create_missing_indexes_concurrently(bind, schema)


def downgrade() -> None:
    bind = _bind()
    schema = _current_schema(bind)

    _drop_indexes_concurrently(bind, schema)

    _require_table(bind, schema, "generation_records")
    if _column_exists(
        bind,
        schema,
        "generation_records",
        _PROMPT_USAGE_COLUMN,
    ):
        op.drop_column(
            "generation_records",
            _PROMPT_USAGE_COLUMN,
            schema=schema,
        )