Zenn
Open2

【Python】alembicでテスト用データベースにもマイグレーションを行う方法

PostgreSQL
Python 3
SQLAlchemy
pytest
alembic
イナジカイナジカ

TL;DR

  • 開発用DBのURLはalembic.iniのsqlalchemy.urlで指定しよう
  • テスト用DBのURLは@pytest.fixtureをつけた関数内で指定しよう

環境

  • Python 3.12
  • fastAPI 0.105.0
  • SQLAlchemy 2.0.23
  • SQLAlchemy Utils 0.41.1
  • alembic 1.13.1
  • pytest 7.4.3
  • psycopg2 2.9.9
  • PostgreSQL 16.1

alembic.iniの設定

開発用DBのURLはalembic.iniのsqlalchemy.iniで指定します。

[alembic]
sqlalchemy.url = postgresql+psycopg2://postgres:postgres@localhost/my_db

環境変数からDBのユーザやパスワード、ホスト名を指定したい場合は以下のようにします。

[alembic]
sqlalchemy.url = postgresql+psycopg2://%(DB_USER)s:%(DB_PASSWORD)s@%(DB_HOST)s/my_db

※%(HOGE)sのように表記します。(環境変数のキーと同一にするとわかりやすいです。)
※alembic.iniが環境変数を読み込むわけではなく、後述のenv.pyで環境変数の値に置換するためのものです。

env.pyの設定

import os # 環境変数参照用

from alembic import context

from sqlalchemy import engine_from_config
from sqlalchemy import pool

# ...[中略]...

config = context.config

# ...[中略]...

def run_migrations_online() -> None:
  # 環境変数から読み込んだ値に置換
  # config.set_section_option("alembic", "alembic.iniで「%(HOGE)s」としたところ", os.environ.get("参照したい環境変数のキー"))
  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"))
  
  # alembic.iniの読み込み
  conf = config.get_section(config.config_ini_section)

  connectable = engine_from_config(
      conf,
      prefix="sqlalchemy.",
      poolclass=pool.NullPool,
  )
  # ...[以下略]...
イナジカイナジカ

テスト用DBセットアップ

テスト実行前にマイグレーションを実行したいので、@pytest.fixtureをつけて定義した関数からマイグレーションを実行できるように定義します。

import os
from typing import Final

import alembic
import alembic.config

import pytest

from sqlalchemy import create_engine
from sqlalchemy.engine import Connection
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.orm.session import close_all_sessions
from sqlalchemy_utils import create_database, drop_database, database_exists

# ① ここでテスト用DBのURLを定義
TEST_DB_URL: Final[str] = "postgresql+psycopg2://postgres:postgres@localhost/test"
"""データベース接続URL"""

# 環境変数から参照する場合は以下のようにしてもかまいません。
# TEST_DB_URL: Final[str] = f"postgresql+psycopg2://{os.environ.get("DB_USER")}:{os.environ.get("DB_PASSWORD")}@{os.environ.get("DB_HOST")}/test"

ALEMBIC_INI_PATH: Final[str] = "alembic.ini"
"""alembic.iniのパス"""

MIGRATIONS_PATH: Final[str] = "src/migrations"
"""マイグレーションスクリプトのあるディレクトリへのパス"""

class TestSession(Session):
  """テスト用DBセッションクラス"""
  def commit(self):
    self.flush()
    self.expire_all()

def migration(connection: Connection):
  """テスト用DBマイグレーション
  
  常に最新版にマイグレーション

  Args:
      connection (Connection): DBアクセスコネクション
  """
  alembic_config = alembic.config.Config(ALEMBIC_INI_PATH) # alembic.iniを読み込む
  alembic_config.set_main_option("script_location", MIGRATIONS_PATH)
  # ② alembic.iniで設定したDBのURLを、①で指定したURLに置換する
  alembic_config.set_main_option("sqlalchemy.url", TEST_DB_URL)
  
  if connection is not None:
    alembic_config.attributes["connection"] = connection
  
  # ③ マイグレーション実行
  alembic.command.upgrade(alembic_config, "head")


@pytest.fixture
def test_database():
  """テスト用DBセットアップ"""
  test_engine = create_engine(
    url=TEST_DB_URL,
    client_encoding="UTF-8",
    echo=True
  )
  """テスト用エンジン"""
  
  if not database_exists(test_engine.url):
    create_database(test_engine.url)

  # テスト用データベースにテーブルを新規作成
  with test_engine.connect() as conn:
    try:
      # マイグレーション
      migration(conn)
      conn.commit()
      
    except Exception as e:
      print(e)
    finally:
      conn.close()

  TestSessionLocal = sessionmaker(
    class_=TestSession,
    bind=test_engine,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False
  )
  """テスト用DB接続セッション"""
  
  yield TestSessionLocal()
  
  # テスト用データベースを削除
  with test_engine.connect() as conn:
    try:
      drop_database(test_engine.url)
      conn.commit()
      
    except Exception as e:
      print(e)
    finally:
      conn.close()
  
  close_all_sessions()
  test_engine.dispose()