FastAPI 개발자가 직접 개발한 FastAPI backend DB 초기화 방법

FastAPI 개발자 tiangolo는 FastAPI를 어떻게 사용할까 궁금해서 찾아본 프로젝트

https://github.com/tiangolo/full-stack-fastapi-postgresql

이번 게시글에서는 DB 초기화를 어떻게 하는지 분석함

소스 개요

이 프로젝트에서는 prestart.sh을 이용하여 DB 초기화를 위한 스크립트 세 개를 수행함

#! /usr/bin/env bash

# Let the DB start
python /app/app/backend_pre_start.py # 데이터베이스 연결과 작동을 확인

# Run migrations
alembic upgrade head # 마이그레이션 수행

# Create initial data in DB
python /app/app/initial_data.py # DB 테이블에 초기 데이터 삽입

---

DB 초기화를 위한 스크립트 설명

backend_pre_start.py

이 스크립트는 데이터베이스가 정상적으로 작동하는지 확인하기 위해 사용될 수 있으며, 특히 Docker 컨테이너와 같은 환경에서 데이터베이스가 준비되었는지 확인할 때 유용함

init() 함수는 데이터베이스에 연결을 시도하고, 연결이 제대로 설정되었는지 확인

데이터베이스 연결이 실패할 경우 재시도할 수 있도록 @retry 데코레이터를 사용하여 init() 함수를 재시도 가능한 함수로 만듦

import logging

from sqlmodel import Session, select
from tenacity import after_log, before_log, retry, stop_after_attempt, wait_fixed

from app.core.db import engine

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

max_tries = 60 * 5  # 5 minutes
wait_seconds = 1


@retry(
    stop=stop_after_attempt(max_tries),
    wait=wait_fixed(wait_seconds),
    before=before_log(logger, logging.INFO),
    after=after_log(logger, logging.WARN),
)
def init() -> None:
    try:
        with Session(engine) as session:
            # Try to create session to check if DB is awake
            session.exec(select(1))
    except Exception as e:
        logger.error(e)
        raise e


def main() -> None:
    logger.info("Initializing service")
    init()
    logger.info("Service finished initializing")


if __name__ == "__main__":
    main()

 

alembic

데이터베이스 마이그레이션 도구로, DB에 테이블을 생성하거나 스키마 변경사항을 적용하는 데 사용

https://bigseok.tistory.com/entry/FastAPI-%EA%B0%9C%EB%B0%9C%EC%9E%90%EA%B0%80-%EC%A7%81%EC%A0%91-%EA%B0%9C%EB%B0%9C%ED%95%9C-FastAPI-backend-alembic

FastAPI 개발자가 직접 개발한 FastAPI backend alembic

initial_data.py

이 스크립트는 데이터베이스 테이블에 초기 데이터를 삽입하는 모듈 init_db를 실행하여 데이터베이스를 초기화

init() 함수는 데이터베이스에 연결한 후, init_db() 함수를 사용하여 초기 데이터를 삽입

import logging

from sqlmodel import Session

from app.core.db import engine, init_db

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def init() -> None:
    with Session(engine) as session:
        init_db(session)


def main() -> None:
    logger.info("Creating initial data")
    init()
    logger.info("Initial data created")


if __name__ == "__main__":
    main()

 

app.core.db.py > init_db

이 스크립트는 데이터베이스 초기화 및 슈퍼 유저 생성을 수행하는데 사용됨

데이터베이스 테이블을 초기화하고 슈퍼 유저를 추가하는 등의 초기 설정을 수행함

설정 파일(app.core.config.settings)에서 가져온 데이터베이스 연결 정보를 사용하여 SQLAlchemy 엔진을 생성하고 애플리케이션과 데이터베이스를 연결

init_db() 함수는 SQLAlchemy의 Session 객체를 인수로 받아 초기화 작업을 수행함

• 설정에서 지정한 슈퍼 유저(superuser)의 이메일을 사용하여 데이터베이스에서 해당 유저를 조회
• 만약 슈퍼 유저가 존재하지 않는다면, 설정에서 지정한 초 슈퍼 유저 정보를 사용하여 UserCreate 모델을 생성
• 생성된 슈퍼 유저 정보를 사용하여 crud.create_user() 함수를 호출하여 초 슈퍼 유저를 데이터베이스에 추가

이 프로젝트는 기본적으로 Alembic 마이그레이션을 사용하여 테이블을 생성하지만,

마이그레이션을 사용하지 않는다면 주석 처리된 코드를 해제하여 SQLAlchemy의 SQLModel.metadata.create_all() 메서드를 주석 해제하여 모든 모델에 대한 테이블을 생성할 수 있도록 함

# init_db
from sqlmodel import Session, create_engine, select

from app import crud
from app.core.config import settings
from app.models import User, UserCreate

engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))


# make sure all SQLModel models are imported (app.models) before initializing DB
# otherwise, SQLModel might fail to initialize relationships properly
# for more details: https://github.com/tiangolo/full-stack-fastapi-postgresql/issues/28


def init_db(session: Session) -> None:
    # Tables should be created with Alembic migrations
    # But if you don't want to use migrations, create
    # the tables un-commenting the next lines
    # from sqlmodel import SQLModel

    # from app.core.engine import engine
    # This works because the models are already imported and registered from app.models
    # SQLModel.metadata.create_all(engine)

    user = session.exec(
        select(User).where(User.email == settings.FIRST_SUPERUSER)
    ).first()
    if not user:
        user_in = UserCreate(
            email=settings.FIRST_SUPERUSER,
            password=settings.FIRST_SUPERUSER_PASSWORD,
            is_superuser=True,
        )
        user = crud.create_user(session=session, user_create=user_in)

 

참고 : FastAPI 개발자가 직접 개발한 FastAPI backend 시리즈

FastAPI 개발자가 직접 개발한 FastAPI backend 프로젝트 구조

FastAPI 개발자가 직접 개발한 FastAPI backend alembic

FastAPI 개발자가 직접 개발한 FastAPI backend DB 초기화 방법

FastAPI 개발자가 직접 개발한 FastAPI backend pyproject.toml

FastAPI 개발자가 직접 개발한 FastAPI backend 설정파일 config.py

FastAPI 개발자가 직접 개발한 FastAPI backend 데이터모델 models.py