SQLModel で設計したモデルを Alembic でマイグレーション

背景

作成したDBモデルを簡単に管理するためにマイグレーション機能を実装しようといくつかその方法を調べたところ、Alembicというものがありました。

Alembicはモデルを再設計した際に、以前までのモデル設計との差分を見て自動でマイグレーションファイルを作成するAuto generate機能というものがあります。これがAlembicの最大の特徴かと思います。

また、AlembicはSQLAlchemyと互換性があるということですが、SQLModelはSQLAlchemyから継承したクラスを利用しているのでSQLModelでもAlembicを使えるんじゃないか？？と思ったので試してみました。

開発環境

• Python 3.10.3

実装

参照記事

こちらの記事に沿って進めていき、必要な部分だけ変更しました。

Alembicをインストール

pip install alembic

Alembicセットアップ

alembic init 任意の環境名
例) alembic init migrations

これらのファイルが自動で作成されます。

$ tree
.
├── alembic.ini
└── migrations
    ├── README
    ├── env.py
    ├── script.py.mako
    └── versions

各ファイルの解説です。（参照記事抜粋）

• alembic.ini

  • alembicのスクリプトが実行される時に読まれる構成ファイル。実行時の設定を記述する。
    • env.py の場所
    • log の出力
    • migration ファイルの命名規則

• env.py

  • マイグレーションツールが実行された時に必ず実行されるPythonスクリプト。
  • SQLAlchemyのEngineを設定や生成を行って、migrationが実行できるようにカスタマイズする。

• script.py.mako

  • 新しいmigrationのスクリプトを生成するために使用される Mako テンプレートファイル。
  • ここにあるものは何でもversions内の新しいファイルを生成するために使用される。

• versions/

  • migrationスクリプトが保存されるディレクトリ

alembic.iniを編集

sqlalchemy.urlの部分を自分のDBの情報に書き換える
DBのユーザー名やパスワードは
マネコン > SecretsManager > 自動作成されているAuroraDBのシークレット > シークレットの値
で確認できます。

例

sqlalchemy.url = mysql+pymysql://ユーザー名:パスワード@127.0.0.1:ローカルのポート/DB名

SQLAlchemyの設定

データベースの設定を行うsettings.pyをトップディレクトリに作成します。

settings.py

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base

path = 'mysql+pymysql://root:@127.0.0.1:3306/DB名'
 
# Engine の作成
Engine = create_engine(
  path,
  encoding="utf-8",
  echo=False
)
Base = declarative_base()

SQLModelのモデルを作成

本来ならここでSQLAlchemyのモデルを作成しますが、データの型検証もできるSQLModelを使う場合で進めていきます。

models.py

from typing import Optional

from sqlmodel import Field, SQLModel

from settings import Base

SQLModel.metadata = Base.metadata


class TestUser(SQLModel, table=True):
    __table_args__ = {'extend_existing': True}
    __tablename__ = "test_users"
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(max_length=32)
    email: str = Field(max_length=64)

ここで注目すべきは、SQLModel.metadata = Base.metadataでSQLModelのメタデータオブジェクトをSQLAlchemyのメタデータオブジェクトにバインドしているという点です。
これによってSQLModelでもAlembicの使用することができるようになります。

alembicの環境設定ファイルであるenv.pyを修正

先ほど設計したモデルクラスをインポートします。
ただし、複数モデルを扱う場合はこちらの記事を参照してenv.pyを修正してください。
(上記引用記事の最後をさらにSQLModel用に
target_metadata = combine_metadata(User.Base.metadata, Group.Base.metadata)
を
target_metadata = combine_metadata(User.metadata, Group.metadata)
に修正してください)

env.py

from models import TestUser # 追加

...

target_metadata = TestUser.metadata # 変更

続けて同じファイルを変更していきます。

env.py

...

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,
    )
    
    # 追加
    url = config.get_main_option("sqlalchemy.url")

    with connectable.connect() as connection:
        context.configure(
            # url追加
            url=url, connection=connection, target_metadata=target_metadata
        )

        with context.begin_transaction():
            context.run_migrations()


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()

テンプレートのscript.py.makoを編集

こちらもSQLModelならではの変更部分です。
下記のように１行だけ追加します。

script.py.mako

"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel # ←追加
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}


def upgrade() -> None:
    ${upgrades if upgrades else "pass"}


def downgrade() -> None:
    ${downgrades if downgrades else "pass"}

マイグレーション実行

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

以下のコマンドでマイグレーションファイルを生成できます。

コマンドと出力

$ alembic revision --autogenerate -m "create tables"
INFO  [alembic.runtime.migration] Context impl MySQLImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.autogenerate.compare] Detected added column 'test_users.status'
  Generating /Users/kawasakiharuki/test/migrations/versions/5e612d3a37ae_create_tables.py ...  done

上記のコマンドを打つとversionsの中にマイグレーションファイルが作成されます。

DBにマイグレーションファイルの内容を反映

コマンドと出力

$ alembic upgrade head
INFO  [alembic.runtime.migration] Context impl MySQLImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade dd91845f6009 -> 5e612d3a37ae, create tables

これでマイグレーションができました。

現在の状態を調べるには以下のコマンドを打ちます。

$ alembic current
INFO  [alembic.runtime.migration] Context impl MySQLImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
5e612d3a37ae (head)

マイグレーション取り消し

以下参照記事抜粋

##
# None -> 8d4b3cb61f6 -> 3f1cb7f09b7a という順序で更新が行われている。
# 現在の状態は3f1cb7f09b7a
##

# 一つ前の状態へ戻す
$ alembic downgrade 8d4b3cb61f6
INFO  [alembic.context] Context impl SQLiteImpl.
INFO  [alembic.context] Will assume transactional DDL.

# 現在の状態。8d4b3cb61f6
$ alembic current
INFO  [alembic.context] Context impl SQLiteImpl.
INFO  [alembic.context] Will assume transactional DDL.
Current revision for sqlite:////home/podhmo/.virtualenvs/pyramid/alembic_sample/data.db: None -> 8d4b3cb61f6, create

# baseを指定すると最初に戻す
$ alembic downgrade base
INFO  [alembic.context] Context impl SQLiteImpl.
INFO  [alembic.context] Will assume transactional DDL.
INFO  [alembic.context] Running downgrade 8d4b3cb61f6 -> None

# 現在の状態。None(最初)
$ alembic current
INFO  [alembic.context] Context impl SQLiteImpl.
INFO  [alembic.context] Will assume transactional DDL.
Current revision for sqlite:////home/podhmo/.virtualenvs/pyramid/alembic_sample/data.db: None

# 元の状態に戻す 3f1cb7f09b7a
$ alembic upgrade 3f1cb7f09b7a
INFO  [alembic.context] Context impl SQLiteImpl.
INFO  [alembic.context] Will assume transactional DDL.
INFO  [alembic.context] Running upgrade None -> 8d4b3cb61f6
INFO  [alembic.context] Running upgrade 8d4b3cb61f6 -> 3f1cb7f09b7a

まとめ

SQLModelで設計したモデルをAlembicでマイグレーションすることができました。
自動でマイグレーションファイルを作成してくれるAuto generateが素晴らしいですね。
ただしAuto generateの欠点としてこちらの記事では以下のように述べられています。

• テーブル名の変更(drop/addになりデータが初期化されます)
• カラム名の変更(drop/addになりデータが初期化されます)
• 名前の付いていないユニーク制約 (変更を検知するにはユニーク制約に名前をつけることは必須(e.g. UniqueConstraint('col1', 'col2', name="my_name"))
• EnumなどのSQLAlchemy特有のカラムの変更
• Columnの順番の指定