mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
feat: Unified response model and exception handling
This commit is contained in:
@@ -0,0 +1,82 @@
|
||||
import json
|
||||
import logging
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.requests import Request
|
||||
|
||||
class ResponseMiddleware(BaseHTTPMiddleware):
|
||||
def __init__(self, app):
|
||||
super().__init__(app)
|
||||
|
||||
async def dispatch(self, request, call_next):
|
||||
response = await call_next(request)
|
||||
|
||||
if isinstance(response, JSONResponse):
|
||||
return response
|
||||
|
||||
if response.headers.get("content-type") == "application/json":
|
||||
try:
|
||||
body = b""
|
||||
async for chunk in response.body_iterator:
|
||||
body += chunk
|
||||
|
||||
raw_data = json.loads(body.decode())
|
||||
|
||||
if isinstance(raw_data, dict) and all(k in raw_data for k in ["code", "data", "msg"]):
|
||||
return JSONResponse(
|
||||
content=raw_data,
|
||||
status_code=response.status_code,
|
||||
headers={
|
||||
k: v for k, v in response.headers.items()
|
||||
if k.lower() not in ("content-length", "content-type")
|
||||
}
|
||||
)
|
||||
|
||||
wrapped_data = {
|
||||
"code": 0,
|
||||
"data": raw_data,
|
||||
"msg": None
|
||||
}
|
||||
|
||||
return JSONResponse(
|
||||
content=wrapped_data,
|
||||
status_code=response.status_code,
|
||||
headers={
|
||||
k: v for k, v in response.headers.items()
|
||||
if k.lower() not in ("content-length", "content-type")
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Response processing error: {str(e)}", exc_info=True)
|
||||
return response
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class exception_handler():
|
||||
@staticmethod
|
||||
async def http_exception_handler(request: Request, exc: HTTPException):
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={
|
||||
"code": exc.status_code,
|
||||
"msg": exc.detail,
|
||||
"data": None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def global_exception_handler(request: Request, exc: Exception):
|
||||
logging.error(f"Error info: {str(exc)}", exc_info=True)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={
|
||||
"code": 500,
|
||||
"msg": str(exc),
|
||||
"data": None
|
||||
}
|
||||
)
|
||||
|
||||
+2
-1
@@ -9,6 +9,7 @@ from starlette.middleware.cors import CORSMiddleware
|
||||
from apps.api import api_router
|
||||
from apps.system.middleware.auth import TokenMiddleware
|
||||
from common.core.config import settings
|
||||
from common.core.response_middleware import ResponseMiddleware
|
||||
|
||||
def custom_generate_unique_id(route: APIRoute) -> str:
|
||||
tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
|
||||
@@ -35,10 +36,10 @@ if settings.all_cors_origins:
|
||||
)
|
||||
|
||||
app.add_middleware(TokenMiddleware)
|
||||
app.add_middleware(ResponseMiddleware)
|
||||
app.include_router(api_router, prefix=settings.API_V1_STR)
|
||||
|
||||
|
||||
|
||||
frontend_dist = os.path.abspath("../frontend/dist")
|
||||
if not os.path.exists(frontend_dist):
|
||||
logging.warning(f"The front-end build directory does not exist: {frontend_dist}")
|
||||
|
||||
Reference in New Issue
Block a user