mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 16:44:15 +08:00
Merge branch 'pr-221' into merge/pr-221-220-219
This commit is contained in:
@@ -682,6 +682,30 @@ def dimension_intraday(
|
||||
# 文件上传
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# 扩展数据 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,
|
||||
@@ -704,9 +728,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":
|
||||
@@ -939,9 +961,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":
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user