[Python] FastAPI 미들웨어로 Sqlalchemy Session 관리하기

 

🌟 들어가며

FastAPI를 사용해 웹 어플리케이션을 구현할 때 Sqlalchemy Session객체의 라이프 사이클을 관리하는 것은 매우 중요합니다. 병행성이 보장되어야하는 환경에서 Session객체는 스레드별로 독립적이지 않기 때문입니다. 따라서 각 스레드별로 독립적으로 Session을 관리해줘야합니다. 이를 위해 미들웨어, ContextVar 그리고 scoped_session을 활용해 병행실행이 필요한 환경해서 Session을 어떻게 관리해야하는지 알아보도록 하겠습니다.

📜 이번 포스팅에서 다루지 않는것

• Sqlalchemy의 Engine, Session의 개념
• FastAPI Depends
• 미들웨어 개념
• ContextVar의 개념

---

🔄 Thread local, Thread Safe?

Sqlalchemy는 BaseModel객체를 다루기 위해 Session을 제공합니다. Session을 통해 DB 레코드를 Python객체로 다룰 수 있고 반대로 Python 객체를 DB 레코드로 변경하는 쿼리를 날려주기도 합니다. 다만 이 Session 객체를 다룰 때에는 주의 해야할 부분이 있습니다.

 

The concurrency model for SQLAlchemy’s Session and AsyncSession is therefore Session per thread, AsyncSession per task.

Session객체는 Thread local(스레드 별로 하나씩 생성되지 않음)하지 않습니다. 우리가 특별한 작업없이 생성하는 모든 변수, 객체들이 Thread local하지 않습니다. 여러 스레드에서 객체를 공유해서 사용할 수 있고, 동시에 해당 객체에 접근할 경우 작업에 문제가 발생할 소지가 있는 Thread Safe 하지 않은 객체인거죠

그렇기 때문에 여러 스레드를 사용하거나, 비동기 작업을 하게될 경우 다른 실행흐름 내에서 공유되던 Session객체를 닫아버리면 다른 작업은 Session 객체를 사용할 수 없게되어버립니다.

이러한 상황에서 Sqlalchemy에서 제안하는 가장 좋은 방법은 ContextManager 패턴을 적용해서 Session이 필요한 순간만 Session을 열고 닫으라는 것 입니다. 다음과 같은 패턴이죠

from sqlalchemy import create_engine
from sqlalchemy.orm import Session

engine = create_engine("postgresql+psycopg2://scott:tiger@localhost/")

# create session and add objects
with Session(engine) as session: # <- Session 열기
    session.add(some_object)
    session.add(some_other_object)
    session.commit()

# <- with 블럭 밖에서는 Session이 닫힘

Sqlalchemy에서 제안하는대로 Context manager를 사용하면 문제가 해결될까요? Session객체를 필요한 순간만 생성하고 닫기 때문에 Thread Safe하게 만들 수 있습니다.

session을 특정 블럭 내부에서 사용하고 다시 ORM 객체의 값을 변경하고 하려고 한다면 어떤일이 일어날까요?

예시 코드를 작성히기 위해 다음과 같이 블로그의 포스팅(Post)과 댓글(Comment)에 관련된 테이블 먼저 정의해보겠습니다.

# Post 모델 정의
class Post(Base):  
    __tablename__ = 'posts'  
  
    id = Column(Integer, primary_key=True, index=True)  
    title = Column(String(255), nullable=False)  # 길이를 지정  
    content = Column(Text, nullable=False)  
  
    comments = relationship('Comment', back_populates='post',  
                            cascade='all, delete-orphan')  
  
# Comment 모델 정의  
class Comment(Base):  
    __tablename__ = 'comments'  
  
    id = Column(Integer, primary_key=True, index=True)  
    content = Column(Text, nullable=False)  
    post_id = Column(Integer, ForeignKey('posts.id'), nullable=False)  
  
    post = relationship('Post', back_populates='comments')

 

그리고 with 블럭 밖에서 값을 변경하는 경우를 예시로 만들어보겠습니다. 다음과 같이 포스팅 정보를 변경하는 FastAPI 라우터를 만들어보겠습니다.

@router.put(  
    "/post/{post_id}",  
)  
def update_post(  
        post_id: int,  
        request: PostSchema = Depends(),  
        session_maker=Depends(get_session_maker),  
):  
    with session_maker() as session:  
        post = session.query(Post).filter(Post.id == post_id).first()  
        post.title = request.title  
        post.content = request.content  
        session.commit()  
  
    # with 블럭 밖에서 session이 close되어서 값 변경이 안됨  
    post.title = "세션 닫힌 후 수정불가능"

 

위 코드에서 post.title의 값은 어떻게 될까요? with 블럭 밖에서는 Session이 유효하지 않기 때문에 ORM 객체의 변경이 DB에 반영되지 않습니다. ORM객체가 DB와 값을 주고받을 수 있는 영역이 with 블럭 내부로 제한되는 것이죠.

또 다른 예시를 한번 들어볼까요? 만일 Post 테이블과 Comment 테이블이 연관관계가 있는 경우 with 블럭 밖에서 연관 객체의 값을 조회하는 경우에 대한 코드를 작성해보겠습니다.

class PostResponseSchema(BaseModel):

    title: str  
    content: str  
    comments: list[str] = []  
  
    @classmethod  
    def mapper(cls, post: Post):  
        return cls(  
            title=post.title,  
            content=post.content,  
            comments=[comment.content for comment in post.comments], # <- 에러 발생
        )
  
@router.get(  
    "/post/{post_id}",  
)  
def get_post(  
        post_id: int,  
        session_maker=Depends(get_session_maker),  
):  
    with session_maker() as session:  
        post = session.query(Post).filter(Post.id == post_id).first()  
  
    return PostResponseSchema.mapper(post)

 

위 코드를 실행하면 PostResponseShchema의 mapper에서 post의 연관 객체인 comments에 접근하려고 할때 DetachedInstanceError 에러가 발생합니다. 닫힌 Session을 사용해서는 연관객체에 접근할 수 없기 때문입니다.

raise orm_exc.DetachedInstanceError(
sqlalchemy.orm.exc.DetachedInstanceError: ... is not bound to a Session; lazy load operation of attribute 'comments' cannot proceed (Background on this error at: https://sqlalche.me/e/20/bhk3)
)

 

---

🔒 Context-local 한 Session

미들웨어를 사용하기 앞서 Context-local 한 Session객체가 무엇인지 알아보겠습니다. Context-local 하다는 의미는 특정 실행 context 에서만 객체가 유효하다는 의미입니다.

Context local에 대해 조금 더 자세하게 설명하기 위해 웹 어플리케이션을 예시로 들어보겠습니다. 웹 어플리케이션에서 Context의 범위는 하나의 사용자 요청입니다. 웹 서비스와 같이 다중 요청이 동시에 처리되는 환경을 가정해보았을 때 Context-local 하게 Session 객체를 사용한다면, 요청마다 독립된 Session 인스턴스를 사용하게 됩니다. 위에서 문제가 되었던, thread-safe 하지 않은 Session 객체 하나를 여러 요청이 공유하는 방식이 아닌, 오직 하나의 컨택스트 내부에서 독립적인 Session 객체를 활용할 수 있게됩니다.

`Context-local` 세션은 Sqlalchemy의 scoped_session을 통해 만들 수 있습니다. scoped_session은 애플리케이션의 현재 스레드 혹은 코루틴의 컨텍스트에 따라 각기 다른 `Session` 객체를 반환합니다.

 

아래는 scoped_session을 사용하는 예시 코드입니다.

from sqlalchemy.orm import scoped_session, sessionmaker
from context_store import get_request_id

SessionFactory = sessionmaker(bind=some_engine)

# request_id에 독립적인 session을 생성
SessionLocal = scoped_session(SessionFactory, scoped_func=get_request_id)

# 사용 예
session = SessionLocal()

---

🛠️  Session 객체와 미들웨어

이번 포스팅에서 보여드릴 FastAPI 미들웨어를 사용해서 Session을 다룰때 필요한 조건입니다.

 

1. Session 객체가 Task 마다 하나씩 만들어져야한다.(Thread-Safe)
2. 특정 로직 흐름 내에서 Session 객체를 열고 닫는 로직이 없어야 한다.
3. 사용자 요청부터 응답까지 유효한 ORM 객체를 사용할 수 있어야 한다.

전체적인 흐름은 다음과 같습니다.

위 흐름을 예시 코드를 작성해 보겠습니다.

 

⚙️ ContextVar, 미들웨어 코드

Context를 구분하기 위해 클라이언트 요청의 `request_id` 를 ContextVar를 사용해 scoped_session 까지 전달하는 코드를 작성해 보겠습니다.

 

1. ContextVar 

요청마다 독립적인 request_id를 관리하는 context 변수를 생성합니다.

# context_store.py
from contextvars import ContextVar  
  
request_id_var: ContextVar[str] = ContextVar("request_id", default="")

 

다음은 reqeust_id를 생성하고, request_id_var에 저장하는 XRequestIdMiddleware 미들웨어입니다. 이 미들웨어를 통해 사용자 요청마다 request_id를 request_id_var이라는 contextVar에 저장합니다.

class XRequestIdMiddleware(BaseHTTPMiddleware):  
    async def dispatch(self, request: Request, call_next):  
        request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))  
        request_id_var.set(request_id)  
        response = await call_next(request)  
        response.headers["X-Request-ID"] = request_id  # 응답에 헤더 추가  
        return response

 

2. scoped session

scoped_session에서는 request_id_var.get 함수를 가져와 scoped_session의 `scopefunc`에 넣어줍니다. 이제 request_id 마다 독립적인 Session을 scoped session으로 부터 얻을 수 있습니다.

# session.py
from sqlalchemy import create_engine  
from sqlalchemy.orm import scoped_session, sessionmaker  
from src.context_store import reqest_id_var

engine = create_engine(settings.database.server_url, pool_pre_ping=True)  
session_maker = sessionmaker(bind=engine)  
session_factory = scoped_session(session_maker, scopefunc=request_id_var.get)

 

3. Session의 라이프 사이클을 관리하는 미들웨어

Session 객체의 라이프 사이클을 관리하는 SessionMiddleware 입니다. 아래 미들웨어에서 Session 객체를 생성하고 요청을 다 처리하고 응답할 때 Session 객체를 반납하게됩니다. 만일 중간에 에러가 발생해 Exception이 발생할 경우 롤백 처리할 수 있게 구현되었습니다.

from session import session_factory

class SessionMiddleware(BaseHTTPMiddleware):  
    async def dispatch(self, request: Request, call_next):  
        session = session_factory()  
        try:  
            response = await call_next(request)
            return response
        except Exception as e:  
            session.rollback()  
            raise e  
        finally:  
            session.close()

 

4. 미들웨어 적용

 

미들웨어를 FastAPI에 적용해 보겠습니다. 커스텀하게 만든 미들웨어는 다음과 같은 방법으로 적용할 수 있습니다.

def middlewares():
    middlewares = []  
  
    # X-Request-ID Middleware  
    middlewares.append(Middleware(XRequestIdMiddleware))  
  
    # Session Middleware  
    middlewares.append(Middleware(SessionMiddleware))  
  
    return middlewares

def create_app():  
    app = FastAPI(  
        title="API Server",  
        version="0.1.0",  
        description="",  
        middleware=middlewares(), # 미들웨어 적용!
        docs_url="/docs",
        redoc="/redoc"  
    )

 

🚀  FastAPI 라우터

마지막으로 Session 라이프사이클을 미들웨어로 관리하는 방법을 적용한 FastAPI로 만든 Router 함수를 만들어보겠습니다.

def get_session():  
    return session_factory()

class PostCreateSchema(BaseModel):  
    title: str  
    content: str  
  
  
class PostResponseSchema(BaseModel):  
    class Config:  
        orm_mode = True  
        title: str  
    content: str  
    comments: list[str] = []  
  
    @classmethod  
    def mapper(cls, post: Post):  
        return cls(  
            title=post.title,  
            content=post.content,  
            comments=[comment.content for comment in post.comments],  
        )  
  
  
@router.post(  
    "/post",
)  
def create_post(  
        request: PostSchema,  
        session: Session = Depends(get_session),  
):  
    post = Post(  
        title=request.title,  
        content=request.content,  
    )  
    session.add(post)  
    session.commit()  
  
    return PostResponseSchema.mapper(post)

@router.get(  
    "/post/{post_id}",  
)  
def get_post(  
        post_id: int,  
        session: Session = Depends(get_session),  
):  
    post = session.query(Post).filter(Post.id == post_id).first()  
  
    return PostResponseSchema.mapper(post)

 

라우터 함수를 살펴보면, Session을 열고 닫는 코드가 없다는 것을 확인할 수 있습니다. 또, 라우터 함수를 벗어나서 PostResponseSchema 클래스의 mapper 함수 내에서도 ORM 객체를 사용할 수 있습니다. 또한 scoped_session을 통해 얻은 Session 객체를 통해 요청 Context마다 유일하고 독립적인 Session 객체를 사용할 수 있습니다.

---

🔚 마치며

지금까지 Middleware, ContextVar, scoped_session을 활용해서 Session 객체의 라이프사이클을 관리하는 방법에 대해 알아봤습니다. 기능 구현을 위해 각 개념을 정말 간단히 소개드렸습니다... 😅 아래 참고했던 블로그 자료를 첨부하니 자세한 내용이 궁금하시면 참고해보세요! 이만 마무리하겠습니다.

 

📚 참고 자료

• [FastAPI] 10. Middleware를 이용한 전후 처리
• FastAPI에서 SQLAlchemy Session 다루는 방법
• Python threading.local 와 ContextVar 비교
• Python contextvar 관련 공식문서