[FastAPI] 데이터베이스 Session 생성을 위한 get_db() 이해

get_db()

FastAPI는 비동기 웹 프레임워크로 데이터베이스와 상호작용을 위해 SQLAlchemy 같은 ORM을 자주 사용한다. 일반적으로 데이터베이스 세션을 관리하기 위해 get_db() 함수를 정의하여 사용하는데, 알아야 할 개념과 함수의 동작 방식에 대해 알아보려고 한다.

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

개념

get_db()에 사용되는 SessionLocal과 yield의 개념을 숙지하기 위해 ORM, Session & SessionLocal, yield vs return을 소개한다. 여기서 설명하는 개념은 Python의 SQLAlchemy를 기준으로 소개하므로 다른 기술에서 사용되는 개념과 차이가 있을 수 있다.

ORM (Object-Relational Mapping)

ORM은 객체 지향 프로그래밍에서 사용되는 Class와 데이터베이스의 Table을 매핑하는 기술이다. SQL이 아닌 Python 코드로 데이터베이스 조작이 가능하도록 해준다.

• Base : SQLAlchemy ORM에서 데이터 모델 (테이블)을 정의할 때 사용하는 기본 클래스
• Base.metadata.create_all(engine) : 정의한 ORM 모델을 기반으로 실제 데이터베이스에 테이블을 생성하는 함수
• sessionmaker : 데이터베이스와 상호 작용하기 위한 세션을 생성하는 객체

from sqlalchemy import Column, Integer, String, create_engine
from sqlalchemy.orm import declarative_base, sessionmaker

DATABASE_URL = "mysql+pymysql://ID:PW@IP:PORT/DBName"
engine = create_engine(DATABASE_URL)

# ORM에서 사용할 Base 클래스 생성
Base = declarative_base()

# User 테이블 정의 (ORM)
class User(Base):
    __tablename__ = "users"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String)

# DB 생성
Base.metadata.create_all(bind=engine)

# 세션 생성
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Session & SessionLocal

• Session
  • 데이터베이스와 연결을 관리하는 객체
  • 트랜잭션을 수행하고, ORM을 통해 데이터베이스와 상호작용할 때 사용
• SessionLocal
  • 데이터베이스의 연결을 유지하면서 요청마다 독립적인 세션을 생성하기 위해 사용
  • 여러 요청이 동시에 처리될 때 충돌을 방지하며, 해당 세션을 재사용하지 않고 요청이 끝나면 닫음

db = SessionLocal() # SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
new_user = User(name="Alice")  # 객체 생성
db.add(new_user)  # 데이터 추가
db.commit()  # 저장
db.refresh(new_user)  # 새로고침
print(new_user.id)  # 자동 생성된 ID 출력
db.close()

yield vs return

return과 yield는 모두 함수에서 값을 반환하는 역할을 하지만, 함수가 실행되는 방식과 메모리 사용 방식에서 차이가 있다.

• return
  • 함수를 즉시 종료하고, 값을 반환
  • 같은 함수를 다시 호출하면, 새롭게 시작

def my_function():
    return "Hello"

print(my_function())  # 출력: Hello
print(my_function())  # 새로 실행됨 (이전 실행과 무관)

• yield
  • 실행을 멈춘 상태 (컨텍스트)를 유지하며, 다음 호출 때 이전 상태 이후부터 실행
  • 모든 데이터를 메모리에 올리지 않고 필요한 순간에 값을 생성해 반환 (메모리 절약)

def my_generator():
    print("첫 번째 실행")
    yield 1
    print("두 번째 실행")
    yield 2
    print("세 번째 실행")
    yield 3

gen = my_generator()  # 제너레이터 객체 생성

print(next(gen))  # 첫 번째 실행 -> 출력: 1
print(next(gen))  # 두 번째 실행 -> 출력: 2
print(next(gen))  # 세 번째 실행 -> 출력: 3

get_db() 동작 과정

get_db()는 return이 아닌 yield를 사용하여 재실행될 경우 finally 구문이 실행되도록 작성돼 있다. db.close()를 적절히

FastAPI에서 Depends(get_db)를 사용하여 의존성을 명시하면 다음과 같이 동작한다.

1. get_db()를 실행하여 yield를 반환해 경로 함수에서 사용
2. 경로 함수 실행이 끝나면, Depens() 내부적으로 get_db.close() 호출 및 finally 구문을 실행

from fastapi import FastAPI, Depends
from sqlalchemy.orm import Session
from .database import SessionLocal

app = FastAPI()

def get_db():
    # 경로 함수 실행 전에 실행됨
    db = SessionLocal()
    try:
        # `db`를 반환하고, 경로 함수에서 사용
        yield db
    finally:
        # 경로 함수 실행이 끝나면 실행됨
        db.close()

@app.get("/")
def read_root(db: Session = Depends(get_db)):
    return {"message": "Hello, FastAPI!"}