alembicのmigrationでDB情報を環境変数から取得できるようにする備忘録

プライベートで作成しているアプリケーションのmigrationツールとしてalembicを利用しています。
migrationのDB情報を環境変数で取得できると便利だなぁと考えてやってみたので備忘録として記事を書きました。

環境の確認と準備

$ python -V
Python 3.9.4
$ pip freeze | grep alembic
alembic==1.6.5

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

今回DBはMariadb 10.5.10を利用します。

alembic.iniの修正

alembic init した際に alembic.ini が作成されて、alembic initファイルにはsqlalchemy.urlセクションが存在します。

sqlalchemy.url = driver://user:pass@localhost/dbname

直書きでも問題ないのであればこのセクションに情報を記述すればよいのですが、今回は環境変数から取得したいので下記のように変更します。
%(DB_USER)s、%(DB_PASSWORD)、%(DB_HOST)sにそれぞれユーザ名、パスワード、ホストの情報が実行時に挿入されます。

sqlalchemy.url = mysql://%(DB_USER)s:%(DB_PASSWORD)s@%(DB_HOST)s/devdatabase

env.pyの修正

alembic init した際に env.py が作成されます。

env.pyの下のほうにrun_migrations_onlineという関数があるのでここを修正していきます。
ドキュメント を参考に修正していきます。
ドキュメントによるとset_section_optionメソッドでalembic.iniで設定した %(変数名)sに対して値を埋め込めるようです。

env.py

config.set_section_option("alembic", "DB_USER", os.environ.get("DB_USER"))
config.set_section_option("alembic", "DB_PASSWORD", os.environ.get("DB_PASSWORD"))
config.set_section_option("alembic", "DB_HOST", os.environ.get("DB_HOST"))

connectable = engine_from_config(
    config.get_section(config.config_ini_section),
    prefix="sqlalchemy.",
    poolclass=pool.NullPool,
)

ちなみに、config.get_section(config.config_ini_section)には下記のように辞書型のデータが入っています。

print(config.get_section(config.config_ini_section))
{'here': '<実行されている場所>', 'script_location': '<alembic.initのscript_locationセクションの値>', 'prepend_sys_path': '.', 'sqlalchemy.url': 'mysql://<user>:<password>@<host>/<db>', 'db_user': '<user>', 'db_password': '<password>', 'db_host': '<host>'}

これで環境変数からDBの情報をalembicに渡すことができるようになりました。
env.pyはpythonコードなので実際には環境変数以外にもファイルなどをパースして埋め込むなどもできると思います。

最後に

alembicのDB情報を環境変数から取得できるようになり自分の中では便利になった感じがあります。
とはいえ、alembic自体最近知って試行錯誤中でありまだまだ改善の余地がありそうです。

参考