mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 17:54:15 +08:00
fix(watchlist): 自选导入上传分块读取, 越限即拒绝而非读完再拒
import-csv / import-image 原先 `await file.read()` 之后才比较长度, 上限只在 整个文件进入内存之后生效, 与 issue #204 修复前的扩展数据上传同类。新增 _read_upload_capped 分块读取, 累计超过上限的那一块立即返回 400 并停止读取, 状态码与文案不变; 内存占用不超过上限 + 一块。 测试: 新增 5 例 (含「越限后不再继续读取」的字节计数断言); test_watchlist_csv 的上传 mock 改为像真实 UploadFile 一样读尽返回 b""。
This commit is contained in:
@@ -40,6 +40,28 @@ _IMPORT_CSV_TYPES = {
|
||||
"text/plain",
|
||||
"application/csv",
|
||||
}
|
||||
# 上传分块读取粒度 (与 ext_data 上传一致)
|
||||
_UPLOAD_CHUNK_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
async def _read_upload_capped(file: UploadFile, max_bytes: int, too_large: str) -> bytes:
|
||||
"""分块读取上传内容, 累计超过 max_bytes 立即拒绝(400), 返回完整字节。
|
||||
|
||||
与 ext_data._write_upload_capped 同类保护: 一次性 `await file.read()` 会先把整个
|
||||
文件读入内存再比较长度, 上限在那之后才生效, 一个远超上限的上传照样把进程内存
|
||||
顶满; 分块读取在越过上限的那一块就停止, 内存占用不超过上限 + 一块。
|
||||
"""
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
while True:
|
||||
chunk = await file.read(_UPLOAD_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
total += len(chunk)
|
||||
if total > max_bytes:
|
||||
raise HTTPException(400, too_large)
|
||||
chunks.append(chunk)
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
class AddRequest(BaseModel):
|
||||
@@ -186,11 +208,9 @@ async def import_from_image(request: Request, file: UploadFile = File(...)):
|
||||
if not ok_type and not ok_ext:
|
||||
raise HTTPException(400, "仅支持 JPG / PNG / WebP / BMP / GIF 图片")
|
||||
|
||||
data = await file.read()
|
||||
data = await _read_upload_capped(file, _MAX_IMPORT_IMAGE_BYTES, "图片过大(上限 12MB)")
|
||||
if not data:
|
||||
raise HTTPException(400, "空文件")
|
||||
if len(data) > _MAX_IMPORT_IMAGE_BYTES:
|
||||
raise HTTPException(400, "图片过大(上限 12MB)")
|
||||
|
||||
existing = {r["symbol"] for r in watchlist.list_symbols()}
|
||||
data_dir = request.app.state.repo.store.data_dir
|
||||
@@ -242,11 +262,9 @@ async def import_from_csv(request: Request, file: UploadFile = File(...)):
|
||||
if not ok_type and not ok_ext:
|
||||
raise HTTPException(400, "仅支持 CSV / TXT 文件")
|
||||
|
||||
data = await file.read()
|
||||
data = await _read_upload_capped(file, _MAX_IMPORT_CSV_BYTES, "文件过大(上限 5MB)")
|
||||
if not data:
|
||||
raise HTTPException(400, "空文件")
|
||||
if len(data) > _MAX_IMPORT_CSV_BYTES:
|
||||
raise HTTPException(400, "文件过大(上限 5MB)")
|
||||
|
||||
data_dir = request.app.state.repo.store.data_dir
|
||||
# 解码与自选/instruments parquet 读取为同步 CPU/IO,挪线程池避免卡事件循环
|
||||
|
||||
@@ -40,7 +40,9 @@ def _mock_upload(*, content: bytes, content_type: str = "text/csv", filename: st
|
||||
file = MagicMock()
|
||||
file.content_type = content_type
|
||||
file.filename = filename
|
||||
file.read = AsyncMock(return_value=content)
|
||||
# 像真实 UploadFile 一样: 第一次 read 返回全部内容, 之后返回 b"" 表示读尽
|
||||
# (端点已改为分块读取, 一直返回同一段内容的 mock 会被当成无限大的文件)
|
||||
file.read = AsyncMock(side_effect=[content, b""])
|
||||
return file
|
||||
|
||||
|
||||
|
||||
@@ -54,7 +54,9 @@ def _mock_upload(
|
||||
file = MagicMock()
|
||||
file.content_type = content_type
|
||||
file.filename = filename
|
||||
file.read = AsyncMock(return_value=content)
|
||||
# 像真实 UploadFile 一样: 第一次 read 返回全部内容, 之后返回 b"" 表示读尽
|
||||
# (端点已改为分块读取, 一直返回同一段内容的 mock 会被当成无限大的文件)
|
||||
file.read = AsyncMock(side_effect=[content, b""])
|
||||
return file
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""自选导入上传体积上限回归测试 (import-csv / import-image)。
|
||||
|
||||
两个端点原先 `data = await file.read()` 之后才比较长度: 上限只在整个文件已经进入
|
||||
内存之后生效, 与 ext_data 上传修复 (issue #204) 前的问题同类。_read_upload_capped
|
||||
分块读取, 越过上限即拒绝(400), 内存占用不超过上限 + 一块。纯逻辑, 不需真实数据源。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from starlette.datastructures import UploadFile
|
||||
|
||||
from app.api import watchlist as watchlist_api
|
||||
from app.api.watchlist import _read_upload_capped
|
||||
|
||||
|
||||
class _CountingStream(io.BytesIO):
|
||||
"""记录被读取的字节数, 用来证明越限后不再继续读。"""
|
||||
|
||||
def __init__(self, data: bytes) -> None:
|
||||
super().__init__(data)
|
||||
self.consumed = 0
|
||||
|
||||
def read(self, size: int = -1) -> bytes: # type: ignore[override]
|
||||
chunk = super().read(size)
|
||||
self.consumed += len(chunk)
|
||||
return chunk
|
||||
|
||||
|
||||
def _upload(stream: io.BytesIO) -> UploadFile:
|
||||
return UploadFile(stream, filename="x.csv")
|
||||
|
||||
|
||||
async def test_under_limit_returns_full_content():
|
||||
data = "code,name\n600519.SH,贵州茅台\n".encode()
|
||||
got = await _read_upload_capped(_upload(io.BytesIO(data)), 1024, "too large")
|
||||
assert got == data
|
||||
|
||||
|
||||
async def test_at_limit_is_allowed():
|
||||
data = b"x" * 64
|
||||
got = await _read_upload_capped(_upload(io.BytesIO(data)), 64, "too large")
|
||||
assert got == data
|
||||
|
||||
|
||||
async def test_over_limit_raises_400_with_the_message_the_endpoint_passes():
|
||||
data = b"x" * 200
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _read_upload_capped(_upload(io.BytesIO(data)), 64, "文件过大")
|
||||
assert exc.value.status_code == 400
|
||||
assert exc.value.detail == "文件过大"
|
||||
|
||||
|
||||
async def test_over_limit_stops_reading_at_the_first_chunk_past_the_cap(monkeypatch):
|
||||
# 4 块的文件, 上限 1.5 块: 第 2 块越限即拒绝, 第 3/4 块不应被读取。
|
||||
monkeypatch.setattr(watchlist_api, "_UPLOAD_CHUNK_BYTES", 16)
|
||||
stream = _CountingStream(b"x" * 64)
|
||||
with pytest.raises(HTTPException):
|
||||
await _read_upload_capped(_upload(stream), 24, "too large")
|
||||
assert stream.consumed == 32
|
||||
|
||||
|
||||
async def test_empty_upload_returns_empty_bytes():
|
||||
got = await _read_upload_capped(_upload(io.BytesIO(b"")), 64, "too large")
|
||||
assert got == b""
|
||||
Reference in New Issue
Block a user