From ebf1a89254b6db9a842eba403e8c6b621219b9c6 Mon Sep 17 00:00:00 2001 From: kevin9327 Date: Sun, 30 Aug 2026 18:51:24 +0900 Subject: [PATCH] =?UTF-8?q?fix(ext-data):=20=E6=89=A9=E5=B1=95=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E4=B8=8A=E4=BC=A0=E5=A2=9E=E5=8A=A0=E4=BD=93=E7=A7=AF?= =?UTF-8?q?=E4=B8=8A=E9=99=90=E5=B9=B6=E5=88=86=E5=9D=97=E8=90=BD=E7=9B=98?= =?UTF-8?q?=20(#204)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 问题: 扩展数据上传 /api/ext-data/{id}/upload 与字段检测 /api/ext-data/detect-fields 都直接 `content = await file.read()`, 无任何字节/行数限制。上传超大 CSV/XLSX 会把整个文件读入内存再交给 Polars/Excel 解析, 可能造成高内存占用、进程 OOM 或 服务不可用。自选截图 OCR 已有 12MB 上限, 扩展上传却没有统一上限。 修复: 抽出 _write_upload_capped(file, dest, max_bytes), 以 1MB 分块把上传写入 临时文件, 累计超过上限即拒绝(413), 不再一次性读入内存; 两个上传入口统一改用它。 上限常量 _MAX_UPLOAD_BYTES = 50MB(对扩展快照数据足够宽裕, 亦便于按需调整)。 兼容性: 后缀白名单、解析与映射逻辑不变; 仅新增体积上限。413 由 detect_fields 的 `except HTTPException: raise` 与 upload_data 的 finally 清理临时目录正常传播。 验证: 新增 tests/test_ext_upload_size_limit.py 覆盖未超限完整写入、恰好等于上限、 超限抛 413 三例, 全过; ruff 无新增告警(B008 与 origin/main 同为 2 个, 均为既有 FastAPI File(...) 默认值写法)。 Co-Authored-By: Claude Opus 4.8 --- backend/app/api/ext_data.py | 32 +++++++++++++--- backend/tests/test_ext_upload_size_limit.py | 41 +++++++++++++++++++++ 2 files changed, 67 insertions(+), 6 deletions(-) create mode 100644 backend/tests/test_ext_upload_size_limit.py diff --git a/backend/app/api/ext_data.py b/backend/app/api/ext_data.py index 6fc005c..48420cb 100644 --- a/backend/app/api/ext_data.py +++ b/backend/app/api/ext_data.py @@ -489,6 +489,30 @@ def dimension_members( # 文件上传 # --------------------------------------------------------------------------- +# 扩展数据 CSV/Excel 上传上限(与自选截图 OCR 的 12MB 上限属同类保护, 见 watchlist.py)。 +# 通过分块写入临时文件, 超限即拒绝, 避免 `await file.read()` 把整个文件读入内存。 +_MAX_UPLOAD_BYTES = 50 * 1024 * 1024 +_UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +async def _write_upload_capped(file: UploadFile, dest: Path, max_bytes: int) -> None: + """分块把上传文件写入 dest, 累计超过 max_bytes 立即拒绝(413)。 + + 避免一次性 `await file.read()` 把整个文件读入内存(大文件可能触发高内存占用、 + 进程 OOM 或服务不可用); 超限时停止继续读取与落盘。 + """ + total = 0 + with dest.open("wb") as f: + while True: + chunk = await file.read(_UPLOAD_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > max_bytes: + raise HTTPException(413, f"文件过大(上限 {max_bytes // (1024 * 1024)}MB)") + f.write(chunk) + + @router.post("/{config_id}/upload") async def upload_data( request: Request, @@ -511,9 +535,7 @@ async def upload_data( tmp_dir = Path(tempfile.mkdtemp()) tmp_path = tmp_dir / f"upload{suffix}" try: - with tmp_path.open("wb") as f: - content = await file.read() - f.write(content) + await _write_upload_capped(file, tmp_path, _MAX_UPLOAD_BYTES) # 直接读取文件,不做列重命名 if suffix == ".csv": @@ -746,9 +768,7 @@ async def detect_fields( tmp_dir = Path(tempfile.mkdtemp()) tmp_path = tmp_dir / f"upload{suffix}" try: - with tmp_path.open("wb") as f: - content = await file.read() - f.write(content) + await _write_upload_capped(file, tmp_path, _MAX_UPLOAD_BYTES) # 直接读取,不要求 symbol 列 if suffix == ".csv": diff --git a/backend/tests/test_ext_upload_size_limit.py b/backend/tests/test_ext_upload_size_limit.py new file mode 100644 index 0000000..c3e14cd --- /dev/null +++ b/backend/tests/test_ext_upload_size_limit.py @@ -0,0 +1,41 @@ +"""扩展数据上传体积上限回归测试(issue #204)。 + +_write_upload_capped 分块把上传写入临时文件, 累计超过上限即拒绝(413), +避免 `await file.read()` 把整个文件读入内存。纯逻辑, 不需真实数据源。 +""" +from __future__ import annotations + +import io + +import pytest +from fastapi import HTTPException +from starlette.datastructures import UploadFile + +from app.api.ext_data import _write_upload_capped + + +def _upload(data: bytes) -> UploadFile: + return UploadFile(io.BytesIO(data), filename="x.csv") + + +async def test_under_limit_writes_full_content(tmp_path): + data = b"code,close\n000001.SZ,10.5\n600000.SH,8.2\n" + dest = tmp_path / "upload.csv" + await _write_upload_capped(_upload(data), dest, max_bytes=1024) + assert dest.read_bytes() == data + + +async def test_at_limit_is_allowed(tmp_path): + data = b"x" * 64 + dest = tmp_path / "upload.csv" + await _write_upload_capped(_upload(data), dest, max_bytes=64) + assert dest.read_bytes() == data + + +async def test_over_limit_raises_413(tmp_path): + data = b"x" * 200 + dest = tmp_path / "upload.csv" + with pytest.raises(HTTPException) as exc: + await _write_upload_capped(_upload(data), dest, max_bytes=64) + assert exc.value.status_code == 413 + assert "过大" in str(exc.value.detail)