FastAPIでアプリ起動時に管理者ユーザを自動生成する

fastapi.FastAPI()のlifespanパラメータを利用して、アプリ起動時および終了時の処理を実行できる。

https://fastapi.tiangolo.com/advanced/events/?h=lifespan#lifespan

これを応用して、アプリ起動時に、DBに存在しなければ管理者ユーザを生成する、という処理を実現可能。

@asynccontextmanager
async def lifespan(app: FastAPI):
    """アプリケーションの起動・終了時に実行される処理を定義するコンテキストマネージャ."""
    try:
        # Adminユーザが存在しない場合は作成する
        is_created = create_super_user()
        if is_created:
            logger.info("Admin user created successfully.")
        else:
            logger.info("Admin user already exists.")
    except Exception:
        logger.exception("Failed to create admin user.")
        raise

    yield


app = FastAPI(lifespan=lifespan)

def create_super_user() -> bool:
    """スーパーユーザを作成する.

    Returns:
        bool: スーパーユーザが作成された場合はTrue, 既に存在する場合はFalse.
    """

    db = next(get_db())

    # ADMIN_EMAILを持つAdminUserが存在する場合はスキップ
    if db.query(AdminUser).filter(AdminUser.email == ADMIN_EMAIL).first():
        return False

    # ユーザー名、メールアドレス、パスワードを指定してAdminUserを作成
    admin_user = AdminUser(
        username=ADMIN_USERNAME,
        email=ADMIN_EMAIL,
        hashed_password=ADMIN_PASSWORD,
    )
    db.add(admin_user)
    db.commit()

    return True

参考までに、get_db()は以下のように実装している。

from typing import Generator

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from libs.constants import DATABASE_URL

IS_POSTGRESQL = "postgresql" in DATABASE_URL

# SQLAlchemy エンジンの作成
engine = create_engine(
    DATABASE_URL,
    connect_args={"check_same_thread": False}
    if DATABASE_URL.startswith("sqlite")
    else {},
)

# セッションの作成
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db() -> Generator[Session, None, None]:
    """SQLAlchemy セッションを取得するジェネレータ関数."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()