修复冻结积分BUG | 拆镜状态异常BUG
This commit is contained in:
+364
@@ -0,0 +1,364 @@
|
||||
"""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,
|
||||
)
|
||||
Reference in New Issue
Block a user