😊

alembicのマイグレーション自動生成をカスタマイズ

2023/02/15に公開

python製のマイグレーションツールのalembicのメモです。

python 3.9.16
alembic 1.9.3
sqlalchemy 2.0.3

https://alembic.sqlalchemy.org/

いろんな使い方ができるうちの1つを発見したのでメモ。

初期の場合において、定義するのがつらい

・共通のカラムなどを同じものを書いていくのがツライ。
・外部キーを設定していくのがツライ。

とりあえず初期の段階でたくさんのテーブルを作る場合はこの辺がつらいです。
運用フェーズになるとそうでもないと思いますが。

・追加のSQLが必要な場合がツライ。

例えばPostgreSQLのRowLevelSecurityを有効化する場合、各テーブルに対してSQLを流すのがツライ。テーブル数にもよりますが手作業だと漏れもでそうでとりあえずツライ。

Modelを定義して自動生成するのが楽

以前チャレンジしたんですが、カラムの並びが自動生成された場合に制御するやり方がわからなかったので断念しましたが、CookBookとしてドキュメントに上がっていたのでやってみました。

https://alembic.sqlalchemy.org/en/latest/cookbook.html#apply-custom-sorting-to-table-columns-within-create-table

env.pyに下記を記載

@writer_sort.rewrites(ops.CreateTableOp)  
def order_columns(context, revision, op):  
    special_names = {  
        "id": -100,  
        "company_id": -90,  
        "deleted": 1000,  
        "create_date": 1003,  
        "update_date": 1004,  
    }  
  
    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 = [  
        col for idx, col in sorted(cols_by_key, key=lambda entry: entry[0])  
    ]  
    re_ops = ops.CreateTableOp(
        table_name=op.table_name,
        columns=columns,
        schema=op.schema,
        _namespace_metadata=op._namespace_metadata,
        _constraints_included=op._constraints_included,
        **op.kw
    )
    re_ops.comment = op.comment # コメントが外れるので追加で指定
    return re_ops

def run_migrations_online() -> None:  
    """Run migrations in 'online' mode.  
  
    In this scenario we need to create an Engine    and associate a connection with the context.  
    """    connectable = engine_from_config(  
        config.get_section(config.config_ini_section),  
        prefix="sqlalchemy.",  
        poolclass=pool.NullPool,  
    )  
  
    with connectable.connect() as connection:  
        context.configure(  
            connection=connection, target_metadata=target_metadata,  
            process_revision_directives=writer,  # order_columns  
        )  
  
        with context.begin_transaction():  
            context.run_migrations()

モデルの書き方

alembic.iniにパスを指定できるようになってます。デフォルトは.

prepend_sys_path = .

ディレクトリ構成は以下で想定

├── alembic.ini
├── migration
│   ├── env.py
│   ├── script.py.mako
│   └── versions
├── models
│   ├── __init__.py
│   ├── base.py
│   ├── mixin.py
│   ├── company.py
│   └── product.py

base.py
書き方は色々あると思います。
https://docs.sqlalchemy.org/en/14/orm/declarative_mixins.html
https://docs.sqlalchemy.org/en/20/orm/declarative_mixins.html

ver.1.4の記述でもver2.0でもOKだった。


# stdlib

# lib
from sqlalchemy import (Boolean, Column, ForeignKey, Integer, MetaData, String,
                        Text, text)
from sqlalchemy.orm import registry
from sqlalchemy.orm.decl_api import DeclarativeMeta
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy_repr import RepresentableBase

# app

convention = {
    "ix": "ix_%(column_0_label)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s"
}

metadata = MetaData(naming_convention=convention)
mapper_registry = registry(metadata=metadata)


class Base(RepresentableBase, metaclass=DeclarativeMeta):
    __abstract__ = True
    registry = mapper_registry
    metadata = mapper_registry.metadata


class RlsModel(Base):
    __abstract__ = True
    id: Mapped[int] = mapped_column(primary_key=True, comment='ID')
    company_id: Mapped[int] = mapped_column(ForeignKey("company.id"), nullable=False, comment='企業ID', )
    deleted: Mapped[bool] = Column(Boolean(), server_default=text('false'), nullable=False, index=True, comment='削除フラグ')

mixin.py

# stdlib
from datetime import datetime
from zoneinfo import ZoneInfo

# lib
from sqlalchemy import (Boolean, Column, Enum, ForeignKey, Integer, String,
                        Text, text, DATE)
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.dialects import postgresql


class TimestampMixin(object):
    create_date: Mapped[datetime] = mapped_column(postgresql.TIMESTAMP(timezone=True),
                                                  default=lambda: datetime.now(ZoneInfo("Asia/Tokyo")),
                                                  comment='作成日')
    update_date: Mapped[datetime] = mapped_column(postgresql.TIMESTAMP(timezone=True),
                                                  nullable=True,
                                                  onupdate=lambda: datetime.now(ZoneInfo("Asia/Tokyo")),
                                                  comment='更新日')

company.py

class Company(Base, TimestampMixin):  
	__tablename__ = 'company'  
    __table_args__ = {'comment': '企業'}  
    id = Column(Integer, primary_key=True, comment='ID')  
    name = Column(String(length=250), nullable=False, comment='名称')  
    is_active = Column(Boolean(), server_default=text('false'), nullable=True, comment='有効フラグ')

product.py

class Product(RlsModel, TimestampMixin):
    __tablename__ = 'product'
    __table_args__ = {'comment': '商品'}
    name = Column(String(length=250), nullable=False, comment='商品名')

CreateTableのあとに追加でSQLを発行したい

例えば、RLSを有効したい場合はこんなのが必要。
せっかく自動生成するなら手作業でいれていきたくない。

op.execute('ALTER TABLE product ENABLE ROW LEVEL SECURITY')

というのが、env.pyに追加することでできた。

https://alembic.sqlalchemy.org/en/latest/api/autogenerate.html#autogen-rewriter

env.py

writer_sort = rewriter.Rewriter()  
writer_rls = rewriter.Rewriter()

@writer_rls.rewrites(ops.CreateTableOp)
def create_rls_permission(context, revision, op):
    rls_op_alter = ops.ExecuteSQLOp(
        sqltext=f"ALTER TABLE {op.table_name} ENABLE ROW LEVEL SECURITY"
    )
    for idx, col in enumerate(op.columns):
        if isinstance(col, Column) and col.key == "company_id":
            return [
                op,
                rls_op_alter,
            ]
    return op

@writer_rls.rewrites(ops.DropTableOp)  
def drop_rls_permission(context, revision, op):  
    rls_op_alter = ops.ExecuteSQLOp(  
        sqltext=f"ALTER TABLE {op.table_name} DISABLE ROW LEVEL SECURITY"  
    )  
    for idx, col in enumerate(op._reverse.columns):  
        if isinstance(col, Column) and col.key == "company_id":  
            return [  
                rls_op_alter,  
                op,  
            ]  
    return op

writer = writer_sort.chain(writer_rls)

ポイントは、rewriterをchainでつなげて、各rewriterにはOperatorを複数返すところ。

生成

alembic revision --autogenerate -m "initial_tables"

結果

こんな感じで、並びもOK、追加のSQLもOKでやりたいことはできました。

def upgrade() -> None:

# ### commands auto generated by Alembic - please adjust! ###

op.create_table('company',

sa.Column('id', sa.Integer(), nullable=False, comment='ID'),

sa.Column('name', sa.String(length=250), nullable=False, comment='名称'),

sa.Column('is_active', sa.Boolean(), server_default=sa.text('false'), nullable=True, comment='有効フラグ'),

sa.Column('create_date', postgresql.TIMESTAMP(timezone=True), nullable=False, comment='作成日'),

sa.Column('update_date', postgresql.TIMESTAMP(timezone=True), nullable=True, comment='更新日'),

sa.PrimaryKeyConstraint('id', name=op.f('pk_company'))

)

op.create_table('product',

sa.Column('id', sa.Integer(), nullable=False, comment='ID'),

sa.Column('company_id', sa.Integer(), nullable=False, comment='企業ID'),

sa.Column('name', sa.String(length=250), nullable=False, comment='商品名'),

sa.Column('create_date', postgresql.TIMESTAMP(timezone=True), nullable=False, comment='作成日'),

sa.Column('update_date', postgresql.TIMESTAMP(timezone=True), nullable=True, comment='更新日'),

sa.ForeignKeyConstraint(['company_id'], ['company.id'], name=op.f('fk_product_company_id_company')),

sa.PrimaryKeyConstraint('id', name=op.f('pk_product'))

)
op.execute('ALTER TABLE product ENABLE ROW LEVEL SECURITY')

# ### end Alembic commands ###

  
  

def downgrade() -> None:

op.execute('ALTER TABLE product DISABLE ROW LEVEL SECURITY')

op.drop_table('product')

op.drop_table('company')

初期のテーブル構築はModel定義してやったほうがテーブルの見直しが入った場合も都度全消しして一括再生性できるので最高です。

Discussion