diff --git a/app/api/endpoints/download.py b/app/api/endpoints/download.py index 53f7249..a07ee9a 100644 --- a/app/api/endpoints/download.py +++ b/app/api/endpoints/download.py @@ -4,7 +4,7 @@ import zipfile import aiofiles import httpx import yaml -from fastapi import APIRouter, Request, Query # 导入FastAPI组件 +from fastapi import APIRouter, Request, Query, HTTPException # 导入FastAPI组件 from starlette.responses import FileResponse from app.api.models.APIResponseModel import ErrorResponseModel # 导入响应模型 @@ -18,7 +18,6 @@ config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.pa with open(config_path, 'r', encoding='utf-8') as file: config = yaml.safe_load(file) - async def fetch_data(url: str, headers: dict = 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' @@ -28,6 +27,23 @@ async def fetch_data(url: str, headers: dict = None): response.raise_for_status() # 确保响应是成功的 return response +# 下载视频专用 +async def fetch_data_stream(url: str, 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: + response.raise_for_status() + + # 流式保存文件 + async with aiofiles.open(file_path, 'wb') as out_file: + async for chunk in response.aiter_bytes(): + await out_file.write(chunk) + + return True @router.get("/download", summary="在线下载抖音|TikTok视频/图片/Online download Douyin|TikTok video/image") async def download_file_hybrid(request: Request, @@ -104,11 +120,18 @@ 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) + # response = await fetch_data(url, headers=__headers) - # 保存文件 - async with aiofiles.open(file_path, 'wb') as out_file: - await out_file.write(response.content) + is_successful = await fetch_data_stream(url, headers=__headers, file_path=file_path) + if not is_successful: + raise HTTPException( + status_code=400, + detail="An error occurred while fetching data" + ) + + # # 保存文件 + # async with aiofiles.open(file_path, 'wb') as out_file: + # await out_file.write(response.content) # 返回文件内容 return FileResponse(path=file_path, filename=file_name, media_type="video/mp4")