mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-11-04 06:15:14 +08:00
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
"""
|
||
2.2.0 在 PostgreSQL 上将历史整数主键的自增机制修复为 Identity。
|
||
|
||
版本标识： 5b3355c964bb
|
||
前置版本： d58298a0879f
|
||
创建时间： 2025-08-19 12:27:08.451371
|
||
"""
|
||
import sqlalchemy as sa
|
||
from alembic import op
|
||
|
||
from app.runtime.log import logger
|
||
from app.runtime.config import settings
|
||
|
||
# Alembic 使用以下标识连接迁移版本链。
|
||
revision = '5b3355c964bb'
|
||
down_revision = 'd58298a0879f'
|
||
branch_labels = None
|
||
depends_on = None
|
||
|
||
|
||
def upgrade() -> None:
|
||
"""
|
||
在 PostgreSQL 上将历史整数主键的自增机制修复为 Identity。
|
||
"""
|
||
if settings.DB_TYPE.lower() == "postgresql":
|
||
# 将SQLite的Sequence转换为PostgreSQL的Identity
|
||
fix_postgresql_sequences()
|
||
|
||
|
||
def fix_postgresql_sequences():
|
||
"""
|
||
遍历 PostgreSQL 公共表，修复历史迁移后整数主键的自增方式。
|
||
"""
|
||
connection = op.get_bind()
|
||
|
||
result = connection.execute(sa.text("""
|
||
SELECT table_name
|
||
FROM information_schema.tables
|
||
WHERE table_schema = 'public'
|
||
AND table_type = 'BASE TABLE'
|
||
"""))
|
||
tables = [row[0] for row in result.fetchall()]
|
||
|
||
logger.info(f"发现 {len(tables)} 个表需要检查序列")
|
||
|
||
for table_name in tables:
|
||
fix_table_sequence(connection, table_name)
|
||
|
||
|
||
def fix_table_sequence(connection, table_name):
|
||
"""
|
||
检查单表的整数主键并按需转换为 Identity。
|
||
|
||
跳过版本表、无 id 列和已有 Identity 的表；发生异常时记录错误并回滚当前连接事务。
|
||
|
||
:param connection: 本次迁移使用的数据库连接，失败时会回滚其当前事务
|
||
:param table_name: 需要检查的数据库表名
|
||
"""
|
||
try:
|
||
# 跳过alembic_version表，它没有id列
|
||
if table_name == 'alembic_version':
|
||
logger.debug(f"跳过表 {table_name}，这是Alembic版本表")
|
||
return
|
||
|
||
result = connection.execute(sa.text(f"""
|
||
SELECT is_identity, column_default
|
||
FROM information_schema.columns
|
||
WHERE table_name = '{table_name}'
|
||
AND column_name = 'id'
|
||
"""))
|
||
|
||
id_column = result.fetchone()
|
||
if not id_column:
|
||
logger.debug(f"表 {table_name} 没有id列，跳过")
|
||
return
|
||
|
||
is_identity, column_default = id_column
|
||
|
||
if is_identity == 'YES' or (column_default and 'GENERATED BY DEFAULT AS IDENTITY' in column_default):
|
||
logger.debug(f"表 {table_name} 的id列已经是Identity类型，跳过")
|
||
return
|
||
|
||
logger.info(f"表 {table_name} 存在序列，需要修复")
|
||
convert_to_identity(connection, table_name)
|
||
|
||
except Exception as e:
|
||
logger.error(f"修复表 {table_name} 序列时出错: {e}")
|
||
# 回滚当前事务，避免影响后续操作
|
||
connection.rollback()
|
||
|
||
|
||
def convert_to_identity(connection, table_name):
|
||
"""
|
||
为整数主键增加 Identity 属性，以当前最大 ID 加一作为起始值。
|
||
|
||
保留主键和索引；仅忽略已存在 Identity 的错误，其他异常继续向上传播。
|
||
|
||
:param connection: 本次迁移使用的数据库连接
|
||
:param table_name: 需要转换主键自增机制的数据库表名
|
||
"""
|
||
try:
|
||
result = connection.execute(sa.text(f"""
|
||
SELECT COALESCE(MAX(id), 0) + 1 as next_value
|
||
FROM "{table_name}"
|
||
"""))
|
||
next_value = result.fetchone()[0]
|
||
|
||
# 直接修改列属性，添加Identity，保持其他约束不变
|
||
# 这种方式不会删除主键约束和索引
|
||
connection.execute(sa.text(f"""
|
||
ALTER TABLE "{table_name}"
|
||
ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY (START WITH {next_value})
|
||
"""))
|
||
|
||
logger.info(f"表 {table_name} 序列已转换为Identity，起始值为 {next_value}")
|
||
|
||
except Exception as e:
|
||
# 如果是已经存在的Identity错误，则忽略
|
||
if "already an identity column" in str(e):
|
||
logger.warn(f"表 {table_name} 的id列已经是Identity类型，忽略此错误: {e}")
|
||
return
|
||
logger.error(f"转换表 {table_name} 序列时出错: {e}")
|
||
raise
|