diff --git a/data_chain/apps/router/document.py b/data_chain/apps/router/document.py index 93178a0a757ba2604875219d57a5e3b8fbca48a8..57211d60a5ed2e9cb0263c456c25f531ab5eaf7d 100644 --- a/data_chain/apps/router/document.py +++ b/data_chain/apps/router/document.py @@ -1,9 +1,10 @@ # Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved. -from fastapi import APIRouter, Depends, Query, Body, File, UploadFile +from fastapi import APIRouter, Depends, Query, Body, File, UploadFile, HTTPException from fastapi.responses import StreamingResponse, HTMLResponse, Response from typing import Annotated import urllib +import time from uuid import UUID from httpx import AsyncClient from typing import Annotated @@ -29,7 +30,9 @@ from data_chain.entities.response_data import ( GetTemporaryDocumentStatusResponse, UploadTemporaryDocumentResponse, GetTemporaryDocumentTextResponse, - DeleteTemporaryDocumentResponse + DeleteTemporaryDocumentResponse, + SpeedTestResult, + SpeedTestResponse ) from data_chain.apps.service.session_service import get_user_sub, verify_user from data_chain.apps.service.router_service import get_route_info @@ -221,3 +224,63 @@ async def delete_temporary_docs( req: Annotated[DeleteTemporaryDocumentRequest, Body()]): doc_ids = await DocumentService.delete_temporary_docs(user_sub, req.ids) return DeleteTemporaryDocumentResponse(result=doc_ids) + + +@router.post('/speed-test', response_model=SpeedTestResponse, dependencies=[Depends(verify_user)]) +async def upload_speed_test( + user_sub: Annotated[str, Depends(get_user_sub)], + file: UploadFile = File(...)): + """ + 网络上传速度测试接口 + - 仅接受单个文件 + - 文件大小限制为最大 50KB + - 文件接收后立即丢弃,不保存到存储系统 + - 返回处理时间用于前端计算上传速度 + """ + # 限制文件大小为 50KB + MAX_FILE_SIZE = 50 * 1024 # 50KB + + # 记录开始时间 + start_time = time.time() + + try: + # 读取文件内容 + content = await file.read() + + # 检查文件大小 + if len(content) > MAX_FILE_SIZE: + raise HTTPException( + status_code=413, + detail=f"文件大小超过限制,最大允许 {MAX_FILE_SIZE // 1024}KB" + ) + + # 记录结束时间 + end_time = time.time() + + # 计算处理时间(毫秒) + processing_time = (end_time - start_time) * 1000 + + # 文件大小(字节) + file_size = len(content) + + # 立即丢弃文件内容,不保存 + del content + + _result = SpeedTestResult( + success=True, + file_size=file_size, + processing_time_ms=processing_time, + ) + return SpeedTestResponse( + result=_result, + message="速度测试完成,文件已丢弃" + ) + + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"速度测试失败: {str(e)}") + finally: + # 确保文件句柄被关闭 + if hasattr(file, 'file') and file.file: + file.file.close() diff --git a/data_chain/entities/response_data.py b/data_chain/entities/response_data.py index cb636eedb43522fe2dee3981184cd4ca44a38c33..0d8d5295841d3694392ab324b52fc3750249e44c 100644 --- a/data_chain/entities/response_data.py +++ b/data_chain/entities/response_data.py @@ -29,6 +29,20 @@ from data_chain.entities.enum import ( from data_chain.parser.parse_result import ParseResult +class SpeedTestResult(BaseModel): + """速度测试结果""" + success: bool = Field(default=False, description="测试是否成功") + file_size: int = Field(default=0, description="文件大小(字节)") + processing_time_ms: float = Field(default=0.0, description="处理时间(毫秒)") + + +class SpeedTestResponse(BaseModel): + """速度测试响应""" + code: int = Field(default=200, description="返回码") + result: SpeedTestResult = Field(default=SpeedTestResult(), description="速度测试结果") + message: str = Field(default="测试失败", description="测试结果消息") + + class ResponseData(BaseModel): """基础返回数据结构"""