FastAPI BaseHTTPMiddleware 관련 이슈

개요

FastAPI에서 Custom Middleware를 생성할 때 일반적으로 BaseHTTPMiddleware를 상속받아서 구현하는 형태로 사용한다. https://www.starlette.io/middleware/#basehttpmiddleware Starlette의 Middleware 관련 문서에 들어가보면 자세한 내용이 나와있다. 하지만 해당 미들웨어의 경우 몇가지 이슈가 존재하는데 그에 관해 간단하게 살펴보고 해결책을 제시하고자 한다.

문제점

첫 번째로 중간에 Request가 끊기면 정상적으로 동작하지 못한다. No response에러가 발생하며 예외가 터지게 된다.

두 번째로 BackgroundTasks와 이슈가 발생한다. 첫 번째 이슈의 연장선으로, Request가 들어오고 BackgroundTask가 실행중일 때 Request가 끊긴다면 BackgroundTask가 정상적으로 동작하지 못한다. 

이에 관해 몇달전부터 많은 이야기가 오고가고 있으며 수정된 PR이 머지되었지만 동일한 현상이 발생하는 등 문제가 지속적으로 발생하고 있다. https://github.com/encode/starlette/issues/1678 에 따르면 BaseHTTPMiddleware를 deprecated시키자는 이야기까지 나오고 있으며 Starlette의 공식 도큐먼트 또한

위처럼 사용을 자제하라고 나와있다.

해결 방법

해결방법은 나름 간단하다. Pure(Raw) ASGI Middleware형태로 수정해주면 된다. 일반적으로 BaseHTTPMiddleware를 상속하여 미들웨어를 생성하는 방법은 아래와 같다.

from starlette.middleware.base import BaseHTTPMiddleware


class CustomMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request, call_next):
        response = await call_next(request)
        return response

위처럼 상속받은 클래스를 하나 생성하고 dispatch() 메소드를 오버라이드하여 필요한 작업을 수행하고 리턴한다. 위 코드를 아래와 같이 수정해준다.

from starlette.types import ASGIApp, Receive, Scope, Send


class CustomMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        await self.app(scope, receive, send)

ASGI 스펙에 맞는 Pure한 Middleware이다. 아무래도 이런 형태로 사용한다면 FastAPI의 request등을 이용하기가 조금 까다롭다. 예를 들어 미들웨어에서 특정 Request/Response에 BackgroundTask를 엮어주고 싶은 경우

background_tasks = BackgroundTasks()
background_tasks.add_task(some_task)
response.background = background_tasks

위와 같이 response객체의 background 속성에 BackgroundTask를 넣어준다면 응답이 나간 이후 해당 태스크가 수행된다. 하지만 Pure ASGI Middleware의 경우 그러한 작업이 까다롭게 된다. (시도해보지 않아서 가능한지조차 모르겠다)

주의 사항

미들웨어를 등록하는 방법이 한가지 더 있다. 

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

위처럼 데코레이터 형태로 등록하는 것인데, 해당 메소드로 들어가보면

def middleware(self, middleware_type: str) -> typing.Callable:
    assert (
        middleware_type == "http"
    ), 'Currently only middleware("http") is supported.'

    def decorator(func: typing.Callable) -> typing.Callable:
        self.add_middleware(BaseHTTPMiddleware, dispatch=func)
        return func

    return decorator

이렇게 middleware_type이 http인 경우, BaseHTTPMiddleware 를 사용하는 형태로 구현되어있다. 따라서 이 부분은 주의가 필요하다.

참고

https://github.com/encode/starlette/pull/1656/files 에 Pure ASGI Middleware 관련 Documentation PR이 올라와있다. 아직 머지되진 않았지만 상세한 설명이 있어 이해에 도움이 될 것이다.