How to sort primary columns to the front in alembic / sqlalchemy

2023-02-20 2-minute read

At work we autogenerated a bunch of sqlalchemy classes from linkml models. Unfortunately, this generated classes had columns in some arbitrary order. The goal was to get at least the primary keys to the front.

First, we had some additional primary keys in a mixin class:


from sqlalchemy import Column
from sqlalchemy.orm import declarative_mixin, declared_attr
from sqlalchemy.sql.sqltypes import BigInteger


@declarative_mixin
class PKMixin:

    @declared_attr
    def add_on_id(self):
        col = Column(
            BigInteger(),
            nullable=False,
            primary_key=True,
        )
        # Moves the column to the front
        col._creation_order = -10
        return col

Using this col._creation_order = -10 trick already pulled the add_on_id column to the front of the table in alembic. The problem is that the main table definition looks like this:


from sqlalchemy import Column, Table
from sqlalchemy.orm import declarative_base
from sqlalchemy.sql.sqltypes import *


class _Base(PKMixin):
    pass


Base = declarative_base(cls=_Base)
metadata = Base.metadata


class TableA(Base):
    __tablename__ = "table_a"
    some_value = Column(Float())
    some_id = Column(Text(), primary_key=True, nullable=False) # AFTER the value!

That resulted in migrations which created the columns in this order:

def upgrade() -> None:
    op.create_table('table_a',
    sa.Column('add_on_id', sa.BigInteger(), nullable=False),
    sa.Column('some_value', sa.Float(), nullable=True),
    sa.Column('some_id', sa.Text(), nullable=False),
    sa.PrimaryKeyConstraint('add_on_id', 'some_id'),
    )

This seems to be because alembic seems to use the column order as in the TableA.__table__.c list, as at that point it is correctly sorted (and the @declared_attr column from the PKMixin at the start!):

    metadata_col_names = OrderedSet(
    c.name for c in metadata_table.c if not c.system
)

The fix was using an additional mixin and some gross hack to sort the column list during class creation time:

from sqlalchemy.orm import declared_attr


class UtilsMixin:
    @declared_attr
    def __mapper_args__(cls):
        # Hack: we want any PK column to end up at the start of the table
        for column in cls.__table__.primary_key:
            # Some columns have already nice numbers, and we want to keep these
            if column._creation_order > 0:
                column._creation_order = -5

        # Horrible, horrible hack, to get primary keys first in alembic
        # Alembic uses the order in __table__.c, not the actual _creation_order
        # __table__.c is an immutable collection with the real data in _collection
        # _collection is a list of tuples with (name, Column(...))
        # So this sorts the thing by the column()._creation_order
        cls.__table__.c._collection.sort(key=lambda k: k[1]._creation_order)

        # Return nothing here as we only mis-used it for our sorting
        return {}


class _Base(UtilsMixin, PKMixin):
    pass

# rest as above

And now we end up with this:

def upgrade() -> None:
    op.create_table('table_a',
    sa.Column('add_on_id', sa.BigInteger(), nullable=False),
    sa.Column('some_id', sa.Text(), nullable=False),
    sa.Column('some_value', sa.Float(), nullable=True),
    sa.PrimaryKeyConstraint('add_on_id', 'some_id'),
    )

Success! :-)