Alembic で create table する migration ファイルを auto gen する際にカラムの順番を調整する

困ったこと

下記のような model を定義した際に...

class Base:
    created_at = Column(DateTime, default=func.now(), nullable=False)
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now(), nullable=False)

BaseModel = declarative_base(metadata=metadata, cls=Base)

class Tenant(BaseModel):
    __tablename__ = "tenant"
    __table_args__ = (UniqueConstraint("name"), {})
    id = Column(String(36), primary_key=True, default=uuid.uuid4)
    name = Column(String(255), nullable=False)

以下のように、ベースに定義したカラムが先にくる migration ファイルが出力されてしまう

op.create_table(
    "tenant",
    sa.Column("created_at", sa.DateTime(), nullable=False),
    sa.Column("updated_at", sa.DateTime(), nullable=False),
    sa.Column("id", sa.String(length=36), nullable=False),
    sa.Column("name", sa.String(length=255), nullable=False),
    sa.PrimaryKeyConstraint("id", name=op.f("pk_tenant")),
    sa.UniqueConstraint("name", name=op.f("uq_tenant_name")),
)

対応方法

migration ファイルを作成するタイミングでリライトする仕組みがあります。
https://alembic.sqlalchemy.org/en/latest/cookbook.html#cookbook-custom-sorting-create-table

ただし、 cook book の通りには動かず、以下のコードに落ち着きました。
エラー AttributeError: 'NoneType' object has no attribute 'schema' が出たのと、
UniqueConstraint が重複して出力される事象が発生しました。

# order_columns の部分でカラムの順番を調整しています。
# UniqueConstraint の columns_combination が重複した場合、一つに絞り込む の部分で UniqueConstraint の重複を排除しています。

@writer.rewrites(ops.CreateTableOp)
def order_columns(_context, _revision, op):
    # order_columns
    special_names = {"id": -100, "created_at": 1001, "updated_at": 1002}
    cols_by_key = [
        (
            special_names.get(col.key, index) if isinstance(col, Column) else 2000,
            col.copy(),
        )
        for index, col in enumerate(op.columns)
    ]
    columns_or_constraints = [
        col for idx, col in sorted(cols_by_key, key=lambda entry: entry[0])
    ]

    # UniqueConstraint の columns_combination が重複した場合、一つに絞り込む
    uc_columns_combinations = []
    for index, column_or_constraint in enumerate(columns_or_constraints):
        if isinstance(column_or_constraint, UniqueConstraint):
            columns_combination = set([col.name for col in column_or_constraint.columns])
            if columns_combination in uc_columns_combinations:
                del columns_or_constraints[index]
            else:
                uc_columns_combinations.append(columns_combination)

    op.columns = columns_or_constraints
    return op

以上です。