core/storage.py

@@ -23,7 +23,7 @@ from core.response import APIResponse
 from core.settings import data_root, settings
 from apps.base.models import FileCodes, UploadChunk
 from core.utils import get_file_url, sanitize_filename
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, StreamingResponse
 
 
 class FileStorageInterface:
@@ -310,20 +310,34 @@ class S3FileStorage(FileStorageInterface):
                     },
                     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",
-                headers={
+                headers=headers
-                    "Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
-                },
             )
+        except HTTPException:
+            raise
         except Exception:
             raise HTTPException(status_code=503, detail="服务代理下载异常，请稍后再试")
 
@@ -602,20 +616,32 @@ class OneDriveFileStorage(FileStorageInterface):
             link = await asyncio.to_thread(
                 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",
-                headers={
+                headers=headers
-                    "Content-Disposition": f'attachment; filename="{filename.encode("utf-8").decode("latin-1")}"'
-                },
             )
+        except HTTPException:
+            raise
         except Exception:
             raise HTTPException(status_code=503, detail="服务代理下载异常，请稍后再试")
 
@@ -776,11 +802,35 @@ class OpenDALFileStorage(FileStorageInterface):
     async def get_file_response(self, file_code: FileCodes):
         try:
             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 = {
-                "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:
             logger.info(e)
@@ -969,26 +1019,32 @@ class WebDAVFileStorage(FileStorageInterface):
         try:
             filename = file_code.prefix + file_code.suffix
             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:
             raise HTTPException(
                 status_code=503, detail=f"WebDAV连接异常: {str(e)}")