feat: 为所有存储后端实现流式文件下载

将文件下载方式从全量读取改为流式传输,减少内存占用并提高大文件下载性能
支持S3、OneDrive、OpenDAL和WebDAV存储后端的流式下载
添加错误处理和HTTP异常捕获
This commit is contained in:
jerryliang
2026-01-16 15:45:12 +08:00
parent f94f819fc1
commit 1313934289
+89 -33
View File
@@ -23,7 +23,7 @@ from core.response import APIResponse
from core.settings import data_root, settings from core.settings import data_root, settings
from apps.base.models import FileCodes, UploadChunk from apps.base.models import FileCodes, UploadChunk
from core.utils import get_file_url, sanitize_filename from core.utils import get_file_url, sanitize_filename
from fastapi.responses import FileResponse from fastapi.responses import FileResponse, StreamingResponse
class FileStorageInterface: class FileStorageInterface:
@@ -310,20 +310,34 @@ class S3FileStorage(FileStorageInterface):
}, },
ExpiresIn=3600, ExpiresIn=3600,
) )
tmp = io.BytesIO()
async def stream_generator():
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.get(link) as resp: async with session.get(link) as resp:
tmp.write(await resp.read()) if resp.status != 200:
tmp.seek(0) raise HTTPException(
content = tmp.read() status_code=resp.status,
tmp.close() detail=f"从S3获取文件失败: {resp.status}"
return Response(
content,
media_type="application/octet-stream",
headers={
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
},
) )
# 设置块大小(例如64KB
chunk_size = 65536
while True:
chunk = await resp.content.read(chunk_size)
if not chunk:
break
yield chunk
from fastapi.responses import StreamingResponse
headers = {
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
}
return StreamingResponse(
stream_generator(),
media_type="application/octet-stream",
headers=headers
)
except HTTPException:
raise
except Exception: except Exception:
raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试") raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试")
@@ -602,20 +616,32 @@ class OneDriveFileStorage(FileStorageInterface):
link = await asyncio.to_thread( link = await asyncio.to_thread(
self._get_file_url, await file_code.get_file_path(), filename self._get_file_url, await file_code.get_file_path(), filename
) )
tmp = io.BytesIO()
async def stream_generator():
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.get(link) as resp: async with session.get(link) as resp:
tmp.write(await resp.read()) if resp.status != 200:
tmp.seek(0) raise HTTPException(
content = tmp.read() status_code=resp.status,
tmp.close() detail=f"从OneDrive获取文件失败: {resp.status}"
return Response(
content,
media_type="application/octet-stream",
headers={
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
},
) )
chunk_size = 65536
while True:
chunk = await resp.content.read(chunk_size)
if not chunk:
break
yield chunk
headers = {
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
}
return StreamingResponse(
stream_generator(),
media_type="application/octet-stream",
headers=headers
)
except HTTPException:
raise
except Exception: except Exception:
raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试") raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试")
@@ -776,12 +802,36 @@ class OpenDALFileStorage(FileStorageInterface):
async def get_file_response(self, file_code: FileCodes): async def get_file_response(self, file_code: FileCodes):
try: try:
filename = file_code.prefix + file_code.suffix filename = file_code.prefix + file_code.suffix
# 尝试使用流式读取器
try:
# OpenDAL 可能提供 reader 方法返回一个异步读取器
reader = await self.operator.reader(await file_code.get_file_path())
except AttributeError:
# 如果 reader 方法不存在,回退到全量读取(兼容旧版本)
content = await self.operator.read(await file_code.get_file_path()) content = await self.operator.read(await file_code.get_file_path())
headers = { headers = {
"Content-Disposition": f'attachment; filename="{filename}"'} "Content-Disposition": f'attachment; filename="{filename}"'
}
return Response( return Response(
content, headers=headers, media_type="application/octet-stream" content, headers=headers, media_type="application/octet-stream"
) )
async def stream_generator():
chunk_size = 65536
while True:
chunk = await reader.read(chunk_size)
if not chunk:
break
yield chunk
headers = {
"Content-Disposition": f'attachment; filename="{filename}"'
}
return StreamingResponse(
stream_generator(),
media_type="application/octet-stream",
headers=headers
)
except Exception as e: except Exception as e:
logger.info(e) logger.info(e)
raise HTTPException(status_code=404, detail="文件已过期删除") raise HTTPException(status_code=404, detail="文件已过期删除")
@@ -969,6 +1019,8 @@ class WebDAVFileStorage(FileStorageInterface):
try: try:
filename = file_code.prefix + file_code.suffix filename = file_code.prefix + file_code.suffix
url = self._build_url(await file_code.get_file_path()) url = self._build_url(await file_code.get_file_path())
async def stream_generator():
async with aiohttp.ClientSession(headers={ async with aiohttp.ClientSession(headers={
"Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}" "Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
}) as session: }) as session:
@@ -978,16 +1030,20 @@ class WebDAVFileStorage(FileStorageInterface):
status_code=resp.status, status_code=resp.status,
detail=f"文件获取失败{resp.status}: {await resp.text()}", detail=f"文件获取失败{resp.status}: {await resp.text()}",
) )
# 读取内容到内存 chunk_size = 65536
content = await resp.read() while True:
return Response( chunk = await resp.content.read(chunk_size)
content=content, if not chunk:
media_type=resp.headers.get( break
"Content-Type", "application/octet-stream" yield chunk
),
headers={ headers = {
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode()}"' "Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode()}"'
}, }
return StreamingResponse(
stream_generator(),
media_type="application/octet-stream",
headers=headers
) )
except aiohttp.ClientError as e: except aiohttp.ClientError as e:
raise HTTPException( raise HTTPException(