From 32755dd86b774e0971a103c8bab95734f16c27e8 Mon Sep 17 00:00:00 2001 From: hadwinfu Date: Sun, 2 Mar 2025 23:18:35 +0800 Subject: [PATCH] =?UTF-8?q?Fix:=20=E8=87=AA=E5=8A=A8=E6=B8=85=E7=90=86?= =?UTF-8?q?=E6=9C=AA=E5=AE=8C=E6=88=90=E7=9A=84=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/api/endpoints/download.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/app/api/endpoints/download.py b/app/api/endpoints/download.py index a07ee9a..740ee43 100644 --- a/app/api/endpoints/download.py +++ b/app/api/endpoints/download.py @@ -28,11 +28,10 @@ async def fetch_data(url: str, headers: dict = None): return response # 下载视频专用 -async def fetch_data_stream(url: str, headers: dict = None, file_path: str = None): +async def fetch_data_stream(url: str, request:Request , headers: dict = None, file_path: str = None): headers = { 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36' } if headers is None else headers.get('headers') - async with httpx.AsyncClient() as client: # 启用流式请求 async with client.stream("GET", url, headers=headers) as response: @@ -41,8 +40,12 @@ async def fetch_data_stream(url: str, headers: dict = None, file_path: str = Non # 流式保存文件 async with aiofiles.open(file_path, 'wb') as out_file: async for chunk in response.aiter_bytes(): + if await request.is_disconnected(): + print("客户端断开连接，清理未完成的文件") + await out_file.close() + os.remove(file_path) + return False await out_file.write(chunk) - return True @router.get("/download", summary="在线下载抖音|TikTok视频/图片/Online download Douyin|TikTok video/image") @@ -122,10 +125,10 @@ async def download_file_hybrid(request: Request, __headers = await HybridCrawler.TikTokWebCrawler.get_tiktok_headers() if platform == 'tiktok' else await HybridCrawler.DouyinWebCrawler.get_douyin_headers() # response = await fetch_data(url, headers=__headers) - is_successful = await fetch_data_stream(url, headers=__headers, file_path=file_path) - if not is_successful: + success = await fetch_data_stream(url, request, headers=__headers, file_path=file_path) + if not success: raise HTTPException( - status_code=400, + status_code=500, detail="An error occurred while fetching data" )