SQLAlchemy自定义AlterColumn编译器未注册致Alembic升级失败

核心原因 #

自动生成的迁移文件中，Alembic误对my_id（IDENTITY列）生成了ALTER COLUMN ... DROP DEFAULT语句，但PostgreSQL要求对IDENTITY列移除默认值时必须使用DROP IDENTITY而非DROP DEFAULT。这是因为SQLAlchemy的autogenerate机制未正确识别IDENTITY列的特殊属性，将其默认值视为普通列的默认值进行对比，从而生成错误的DDL语句。

解决方案 #

1. 临时修复：手动修改迁移文件 #

检查自动生成的迁移文件（如versions/xxxxxx_add_authors_column.py），找到针对my_id列的ALTER TABLE ... ALTER COLUMN my_id DROP DEFAULT语句，直接删除该行即可。因为IDENTITY列的默认值由数据库自动维护，无需手动调整。

2. 永久修复：让SQLAlchemy正确识别IDENTITY列 #

调整模型定义（推荐）

替换原有的自定义CreateColumn编译逻辑，直接使用SQLAlchemy原生的Identity类型来定义自增主键，这样SQLAlchemy能自动识别IDENTITY列，避免autogenerate误判：

from sqlalchemy import Column, Integer, String, ARRAY
from sqlalchemy.dialects.postgresql import Identity

class YourTable(Base):
    __tablename__ = 'your_table'
    my_id = Column(Integer, Identity(always=False), primary_key=True)
    # 其他列...
    authors = Column(ARRAY(String(128)), nullable=False, server_default="ARRAY[]::varchar[]")

这种方式无需自定义编译逻辑，SQLAlchemy会自动生成正确的IDENTITY语句，autogenerate时也不会对该列生成错误的DROP DEFAULT操作。

修复自定义编译逻辑（若需保留原有编译）

如果必须保留自定义的CreateColumn编译，需要给列添加标识以便AlterColumn编译时识别：

from sqlalchemy.sql.ddl import CreateColumn, AlterColumn
from sqlalchemy.ext.compiler import compiles

@compiles(CreateColumn, "postgres")
def compile_create_column(element, compiler, **kwargs):
    if isinstance(element.type, Integer) and element.primary_key:
        # 标记该列为IDENTITY列
        element.column.info['is_identity'] = True
        return f"{element.column.name} INT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY"
    return compiler.visit_create_column(element, **kwargs)

@compiles(AlterColumn, "postgres")
def compile_alter_column(element, compiler, **kwargs):
    # 判断是否为IDENTITY列的DROP DEFAULT操作
    is_drop_default = element.modifications.get('default') is None
    is_identity_col = element.column.info.get('is_identity', False)
    
    if is_drop_default and is_identity_col:
        table_name = compiler.preparer.format_table(element.table)
        col_name = compiler.preparer.format_column(element.column)
        return f"ALTER TABLE {table_name} ALTER COLUMN {col_name} DROP IDENTITY"
    
    # 其他情况使用默认编译逻辑
    return compiler.visit_alter_column(element, **kwargs)

注意：确保在Alembic的env.py中导入这些编译函数，否则运行时不会生效。

3. 排查编译扩展未生效的原因 #

• 确认自定义编译函数已正确导入到Alembic的env.py中，Alembic运行时需要加载这些扩展才能生效。
• 检查AlterColumn编译函数的逻辑是否正确匹配到目标语句：比如是否正确判断了element.modifications中的default属性，以及列的is_identity标记是否存在。

内容的提问来源于stack exchange，提问作者Luigi D.