AiOpsWechat/app/middleware.py

import json
import time
from fastapi import FastAPI, Request,HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from pydantic import BaseModel
from datetime import datetime
import logging


from common.log import logger




class Result(BaseModel):
    code: int
    message: str
    status: str
    
class ResponseData(BaseModel):
    data: dict|list|None
    result: Result
    timestamp: str




async def http_context(request: Request, call_next):
    # 记录请求信息
    request_info = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "method": request.method,
        "url": str(request.url),
        "body": await request.body() if request.method in ["POST", "PUT", "PATCH"] else None,
    }
    logger.info(f"请求: {json.dumps(request_info, separators=(',', ':'), default=str, ensure_ascii=False)}")
    
    # 调用下一个中间件或路由处理程序
    response = await call_next(request)
    
    # 如果响应状态码为 200，则格式化响应为统一格式
    if response.status_code == 200:
        try:
            response_body = b""
            async for chunk in response.body_iterator:
                response_body += chunk
            response_body_str = response_body.decode("utf-8")
            business_data = json.loads(response_body_str)
        except Exception as e:
            business_data = {"error": f"Unable to decode response body: {str(e)}"}
        
        if "code" in business_data:
            message=business_data.get("message","请求失败！")
            result = ResponseData(
                data=None,
                result=Result(code=business_data.get("code",500), message=message, status="failed"),
                timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
            )
        else:
            # 构造统一格式的响应
            result = ResponseData(
                data=business_data,
                result=Result(code=200, message="请求成功！", status="succeed"),
                timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
            )
        response_info = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "status_code": response.status_code,
            "headers": dict(response.headers),
            "body": result.dict(),
        }
        logger.info(f"响应: {json.dumps(response_info, separators=(',', ':'), default=str, ensure_ascii=False)}")
        
        # 返回修改后的响应
        return JSONResponse(content=result.model_dump())

    else:
        print(response)
        print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
        message = "请求失败！"
        # 如果响应状态码不为 200，则记录响应信息
        try:
            response_body = b""
            async for chunk in response.body_iterator:
                response_body += chunk
            response_body_str = response_body.decode("utf-8")
            business_data = json.loads(response_body_str)
        except Exception as e:
            business_data = {"error": f"Unable to decode response body: {str(e)}"}
            # 根据不同状态码定制 message 字段
        
            if response.status_code == 404:
                message = e.detail
            elif response.status_code == 400:
                message = "请求参数错误"
            elif response.status_code == 500:
                message = "服务器内部错误"
        

        # 你可以根据不同的状态码设置更详细的错误消息
        
        # 构造统一格式的响应
        result = ResponseData(
            data={},  # 返回空数据
            result=Result(
                code=response.status_code, 
                message=message,  # 根据状态码返回详细信息
                status="failed"
            ),
            timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
        )
        response_info = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "status_code": response.status_code,
            "headers": dict(response.headers),
            "body": result.dict(),
        }
        logger.info(f"响应: {json.dumps(response_info, separators=(',', ':'), default=str, ensure_ascii=False)}")
        
        # 返回修改后的响应
        return JSONResponse(content=result.model_dump(), status_code=response.status_code)
    
async def http_context_2(request: Request, call_next):
    # 记录请求信息
    request_body = None
    if request.method in ["POST", "PUT", "PATCH"]:
        try:
            request_body = await request.json()  # 使用 .json()，避免影响 FastAPI 解析
        except Exception:
            request_body = "无法解析 JSON"
    
    request_info = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "method": request.method,
        "url": str(request.url),
        "body": request_body,
    }
    logger.info(f"请求: {json.dumps(request_info, separators=(',', ':'), ensure_ascii=False)}")

    # 继续处理请求
    response = await call_next(request)

    # 如果是 422 错误，直接返回，避免 Pydantic 解析错误
    if response.status_code == 422:
        return response

    # 处理正常请求
    try:
        response_body = b""
        async for chunk in response.body_iterator:
            response_body += chunk
        response_body_str = response_body.decode("utf-8")
        business_data = json.loads(response_body_str)
    except Exception as e:
        business_data = {"error": f"无法解析响应体: {str(e)}"}

    if "code" in business_data:
        message = business_data.get("message", "请求失败！")
        result = ResponseData(
            data=None,
            result=Result(code=business_data.get("code", 500), message=message, status="failed"),
            timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
        )
    else:
        result = ResponseData(
            data=business_data,
            result=Result(code=200, message="请求成功！", status="succeed"),
            timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
        )

    response_info = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "status_code": response.status_code,
        "headers": dict(response.headers),
        "body": result.dict(),
    }
    logger.info(f"响应: {json.dumps(response_info, separators=(',', ':'), ensure_ascii=False)}")

    return JSONResponse(content=result.model_dump())