71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
import json
|
|
import time
|
|
from typing import Callable
|
|
|
|
from fastapi import APIRouter, Request, Response
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.routing import APIRoute
|
|
|
|
from config.setting import APP_VERSION
|
|
|
|
|
|
class TimedRoute(APIRoute):
|
|
def get_route_handler(self) -> Callable:
|
|
original_route_handler = super().get_route_handler()
|
|
|
|
async def custom_route_handler(request: Request) -> Response:
|
|
start = time.perf_counter()
|
|
response: Response = await original_route_handler(request)
|
|
duration = time.perf_counter() - start
|
|
|
|
# JSON 응답만 처리
|
|
if getattr(response, "media_type", None) == "application/json":
|
|
body_bytes = b""
|
|
|
|
# 1) 스트리밍 응답이면 모두 수집
|
|
body_iter = getattr(response, "body_iterator", None)
|
|
if body_iter is not None:
|
|
async for chunk in body_iter:
|
|
body_bytes += chunk
|
|
else:
|
|
# 일반 응답
|
|
body_bytes = getattr(response, "body", b"")
|
|
|
|
# 2) 파싱 시도
|
|
try:
|
|
text = body_bytes.decode("utf-8") if body_bytes else ""
|
|
payload = json.loads(text) if text else None
|
|
except Exception:
|
|
# 파싱 실패: 바디/길이 불일치 위험 없도록 원본 그대로 반환
|
|
return response
|
|
|
|
# 3) dict일 때만 필드 주입
|
|
if isinstance(payload, dict):
|
|
payload.setdefault("app_version", APP_VERSION)
|
|
payload.setdefault("process_time", f"{duration:.4f}")
|
|
|
|
# 4) 새 JSONResponse로 재구성 (Content-Length 자동 일치)
|
|
# 원본 상태코드/헤더/미디어타입 유지
|
|
new_headers = dict(response.headers)
|
|
# 압축/길이 관련 헤더는 제거(재계산되도록)
|
|
for h in ("content-length", "Content-Length", "content-encoding", "Content-Encoding"):
|
|
new_headers.pop(h, None)
|
|
|
|
return JSONResponse(
|
|
content=payload,
|
|
status_code=response.status_code,
|
|
headers=new_headers,
|
|
media_type=response.media_type,
|
|
)
|
|
|
|
# JSON 아니라면 바디 불문, 원본 그대로
|
|
return response
|
|
|
|
return custom_route_handler
|
|
|
|
|
|
class CustomAPIRouter(APIRouter):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.route_class = TimedRoute
|