Closed2
【Python】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()
このスクラップは22日前にクローズされました