alembic + SQLModelでデータベースをマイグレーションする

バックエンドをPython + MySQL(PlanetScale) + SQLModel で開発しており、データベースマイグレーションにはalembicを使っています。alembicは、SQLAlchemy の開発者である Mike Bayer さんが作ったツールです。alembicを使う手順をまとめます。

alembicインストール

poetry add alembic --group dev
alembic init alembic

alembic設定ファイルをSQLModel向けに修正

設定ファイル/alembic/env.pyと、自動生成されるスクリプトのテンプレートファイル/alembic/script.py.mako を修正します。

env.py

SQLModelのインポートを追加します。

from sqlmodel import SQLModel

21行目あたりのtarget_metadataをSQLModelに切り替えます。

target_metadata = None -> target_metadata = SQLModel.metadata

alembic/script.py.mako

SQLModelのインポートを追加します。

import sqlmodel

接続先データベースの変更

設定ファイル /alembic.ini 64行目あたりに書かれている接続先データベースを変更します。

• ローカル開発環境 sqlalchemy.url = sqlite:///test.db

• MySQL(PlanetScale) sqlalchemy.url = mysql+pymysql://127.0.0.1:3306/dbname

新規テーブルの登録

新しくテーブルを作りたい場合は、env.pyにSQLModelテーブルクラスのインポートを追加します。

from common.model.hoge import Hoge
__all__ = ["Hoge"]

マイグレーションファイルの自動生成

versions/の中にマイグレーションファイルが生成されます。自動生成では対応できていない場合があるので、必ずファイルの中身を確認します。

alembic revision --autogenerate -m "add column fuga"
alembic upgrade head

マイグレーションファイルを修正する必要のあるケースをいくつか紹介します。

NOT NULL制約をつけたカラムを追加する場合

デフォルト値を設定しないと、マイグレーションを実施する際エラーになります。

ERROR: op.add_column("hoge", sa.Column("updated_at", sa.DateTime(), nullable=False))

OK: op.add_column("hoge", sa.Column("updated_at", sa.DateTime(), nullable=False, server_default="2023-08-24 15:26:19.220144"))

カラムをリネームする場合

自動生成されたコードはdrop_columnしてからadd_columnするコードになっています。そのまま実行するとデータが消えてしまうので、alter_columnに修正します。

with op.batch_alter_table("hoge", schema=None) as batch_op:
    batch_op.alter_column("pub_date", nullable=False, new_column_name="pub_datetime", existing_type=sqlalchemy.DateTime)

SQLiteでカラムをリネームする場合

SQLiteにはalter columnがありません。そこで、render_as_batch=Trueを指定するとSQLiteでも実行できるようになります。振る舞いとしては、新しいテーブルに新しいカラムを作成して、データを一方からもう一方にコピーします。

with connectable.connect() as connection:
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        render_as_batch=True,
    )

なお、バッチモードはSQLite以外のデータベースには影響しません。通常のマイグレーションが実行されます。 https://alembic.sqlalchemy.org/en/latest/batch.html#batch-mode-with-autogenerate

This mode is safe to use in all cases, as the Operations.batch_alter_table() directive by default only takes place for SQLite; other backends will behave just as they normally do in the absence of the batch directives.