feat: 为所有存储后端实现流式文件下载
将文件下载方式从全量读取改为流式传输,减少内存占用并提高大文件下载性能 支持S3、OneDrive、OpenDAL和WebDAV存储后端的流式下载 添加错误处理和HTTP异常捕获
This commit is contained in:
+105
-49
@@ -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 with aiohttp.ClientSession() as session:
|
async def stream_generator():
|
||||||
async with session.get(link) as resp:
|
async with aiohttp.ClientSession() as session:
|
||||||
tmp.write(await resp.read())
|
async with session.get(link) as resp:
|
||||||
tmp.seek(0)
|
if resp.status != 200:
|
||||||
content = tmp.read()
|
raise HTTPException(
|
||||||
tmp.close()
|
status_code=resp.status,
|
||||||
return Response(
|
detail=f"从S3获取文件失败: {resp.status}"
|
||||||
content,
|
)
|
||||||
|
# 设置块大小(例如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",
|
media_type="application/octet-stream",
|
||||||
headers={
|
headers=headers
|
||||||
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
|
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 with aiohttp.ClientSession() as session:
|
async def stream_generator():
|
||||||
async with session.get(link) as resp:
|
async with aiohttp.ClientSession() as session:
|
||||||
tmp.write(await resp.read())
|
async with session.get(link) as resp:
|
||||||
tmp.seek(0)
|
if resp.status != 200:
|
||||||
content = tmp.read()
|
raise HTTPException(
|
||||||
tmp.close()
|
status_code=resp.status,
|
||||||
return Response(
|
detail=f"从OneDrive获取文件失败: {resp.status}"
|
||||||
content,
|
)
|
||||||
|
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",
|
media_type="application/octet-stream",
|
||||||
headers={
|
headers=headers
|
||||||
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
|
except HTTPException:
|
||||||
|
raise
|
||||||
except Exception:
|
except Exception:
|
||||||
raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试")
|
raise HTTPException(status_code=503, detail="服务代理下载异常,请稍后再试")
|
||||||
|
|
||||||
@@ -776,11 +802,35 @@ 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
|
||||||
content = await self.operator.read(await file_code.get_file_path())
|
# 尝试使用流式读取器
|
||||||
|
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())
|
||||||
|
headers = {
|
||||||
|
"Content-Disposition": f'attachment; filename="{filename}"'
|
||||||
|
}
|
||||||
|
return Response(
|
||||||
|
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 = {
|
headers = {
|
||||||
"Content-Disposition": f'attachment; filename="{filename}"'}
|
"Content-Disposition": f'attachment; filename="{filename}"'
|
||||||
return Response(
|
}
|
||||||
content, headers=headers, media_type="application/octet-stream"
|
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)
|
||||||
@@ -969,26 +1019,32 @@ 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 with aiohttp.ClientSession(headers={
|
|
||||||
"Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
|
async def stream_generator():
|
||||||
}) as session:
|
async with aiohttp.ClientSession(headers={
|
||||||
async with session.get(url) as resp:
|
"Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
|
||||||
if resp.status != 200:
|
}) as session:
|
||||||
raise HTTPException(
|
async with session.get(url) as resp:
|
||||||
status_code=resp.status,
|
if resp.status != 200:
|
||||||
detail=f"文件获取失败{resp.status}: {await resp.text()}",
|
raise HTTPException(
|
||||||
)
|
status_code=resp.status,
|
||||||
# 读取内容到内存
|
detail=f"文件获取失败{resp.status}: {await resp.text()}",
|
||||||
content = await resp.read()
|
)
|
||||||
return Response(
|
chunk_size = 65536
|
||||||
content=content,
|
while True:
|
||||||
media_type=resp.headers.get(
|
chunk = await resp.content.read(chunk_size)
|
||||||
"Content-Type", "application/octet-stream"
|
if not chunk:
|
||||||
),
|
break
|
||||||
headers={
|
yield chunk
|
||||||
"Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode()}"'
|
|
||||||
},
|
headers = {
|
||||||
)
|
"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(
|
||||||
status_code=503, detail=f"WebDAV连接异常: {str(e)}")
|
status_code=503, detail=f"WebDAV连接异常: {str(e)}")
|
||||||
|
|||||||
Reference in New Issue
Block a user