117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
"""数据库修改
|
||
|
||
Revision ID: baf54bec4cb5
|
||
Revises: 6cf2bcc8915f
|
||
Create Date: 2026-05-28 13:40:31.609405
|
||
"""
|
||
from typing import Sequence, Union
|
||
|
||
from alembic import op
|
||
import sqlalchemy as sa
|
||
|
||
|
||
# revision identifiers, used by Alembic.
|
||
revision: str = 'baf54bec4cb5'
|
||
down_revision: Union[str, None] = '6cf2bcc8915f'
|
||
branch_labels: Union[str, Sequence[str], None] = None
|
||
depends_on: Union[str, Sequence[str], None] = None
|
||
|
||
|
||
TABLE_NAME = "credit_ratios"
|
||
|
||
|
||
def _has_index(table_name: str, index_name: str) -> bool:
|
||
bind = op.get_bind()
|
||
inspector = sa.inspect(bind)
|
||
indexes = inspector.get_indexes(table_name)
|
||
return any(index.get("name") == index_name for index in indexes)
|
||
|
||
|
||
def _has_foreign_key(table_name: str, fk_name: str) -> bool:
|
||
bind = op.get_bind()
|
||
inspector = sa.inspect(bind)
|
||
foreign_keys = inspector.get_foreign_keys(table_name)
|
||
return any(fk.get("name") == fk_name for fk in foreign_keys)
|
||
|
||
|
||
def _create_index_if_not_exists(
|
||
index_name: str,
|
||
table_name: str,
|
||
columns: list[str],
|
||
unique: bool = False,
|
||
) -> None:
|
||
if not _has_index(table_name, index_name):
|
||
op.create_index(index_name, table_name, columns, unique=unique)
|
||
|
||
|
||
def _drop_index_if_exists(index_name: str, table_name: str) -> None:
|
||
if _has_index(table_name, index_name):
|
||
op.drop_index(index_name, table_name=table_name)
|
||
|
||
|
||
def upgrade() -> None:
|
||
# 索引存在才不重复创建，避免 DuplicateTable
|
||
_create_index_if_not_exists(
|
||
"ix_credit_ratios_gen_type",
|
||
TABLE_NAME,
|
||
["gen_type"],
|
||
unique=False,
|
||
)
|
||
|
||
_create_index_if_not_exists(
|
||
"ix_credit_ratios_gen_type_engine_resolution",
|
||
TABLE_NAME,
|
||
["gen_type", "model_config_id", "resolution"],
|
||
unique=False,
|
||
)
|
||
|
||
_create_index_if_not_exists(
|
||
"ix_credit_ratios_gen_type_resolution",
|
||
TABLE_NAME,
|
||
["gen_type", "resolution"],
|
||
unique=False,
|
||
)
|
||
|
||
_create_index_if_not_exists(
|
||
"ix_credit_ratios_model_config_id",
|
||
TABLE_NAME,
|
||
["model_config_id"],
|
||
unique=False,
|
||
)
|
||
|
||
_create_index_if_not_exists(
|
||
"ix_credit_ratios_resolution",
|
||
TABLE_NAME,
|
||
["resolution"],
|
||
unique=False,
|
||
)
|
||
|
||
# 外键存在才删除，避免 UndefinedObject
|
||
if _has_foreign_key(TABLE_NAME, "credit_ratios_model_config_id_fkey"):
|
||
op.drop_constraint(
|
||
"credit_ratios_model_config_id_fkey",
|
||
TABLE_NAME,
|
||
type_="foreignkey",
|
||
)
|
||
|
||
|
||
def downgrade() -> None:
|
||
# 注意：
|
||
# 如果 model_config_id 已经改成 ImageEngine / VideoEngine 的 id，
|
||
# 恢复到 model_configs 外键可能因为历史数据不匹配而失败。
|
||
# 当前业务是删除外键，所以 downgrade 这里只做尽量安全处理。
|
||
|
||
if not _has_foreign_key(TABLE_NAME, "credit_ratios_model_config_id_fkey"):
|
||
op.create_foreign_key(
|
||
"credit_ratios_model_config_id_fkey",
|
||
TABLE_NAME,
|
||
"model_configs",
|
||
["model_config_id"],
|
||
["id"],
|
||
)
|
||
|
||
_drop_index_if_exists("ix_credit_ratios_resolution", TABLE_NAME)
|
||
_drop_index_if_exists("ix_credit_ratios_model_config_id", TABLE_NAME)
|
||
_drop_index_if_exists("ix_credit_ratios_gen_type_resolution", TABLE_NAME)
|
||
_drop_index_if_exists("ix_credit_ratios_gen_type_engine_resolution", TABLE_NAME)
|
||
_drop_index_if_exists("ix_credit_ratios_gen_type", TABLE_NAME) |