mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 23:44:16 +08:00
问题: 扩展数据上传 /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 <noreply@anthropic.com>
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
"""扩展数据上传体积上限回归测试(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)
|