Stream uploaded attachments to disk and cover the progress callback

UploadedFile already holds the whole attachment in memory; getvalue()
made a second full copy of each of the 230 files before writing. Copy
through the file interface instead. Smoke tests pin the progress
callback contract: once per XML file, including unreadable ones.

Refs #1

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Andrey Epifancev 2026-10-09 19:14:43 +04:00
parent 6668858b6f
commit 27a1f22be4
3 changed files with 116 additions and 1 deletions

View File

@ -1,4 +1,5 @@
import asyncio
import shutil
import tempfile
from pathlib import Path
@ -122,7 +123,12 @@ def _save_attachments(uploaded_files, target_dir: Path) -> list[Path]:
holder = target_dir / str(index)
holder.mkdir(parents=True, exist_ok=True)
path = holder / uploaded.name
path.write_bytes(uploaded.getvalue())
# Потоковое копирование: UploadedFile уже целиком лежит в памяти
# Streamlit, а getvalue() делал бы вторую полную копию каждого
# вложения перед записью на диск.
uploaded.seek(0)
with path.open("wb") as target:
shutil.copyfileobj(uploaded, target)
paths.append(path)
return paths

View File

@ -0,0 +1,59 @@
from __future__ import annotations
import zipfile
from pathlib import Path
from app.pipeline import service
from app.pipeline.service import run_batch
from tests.smoke.test_batch_smoke import _patch_opencloud, _valid_xml
def test_progress_callback_is_called_once_per_xml_file(
tmp_path: Path, monkeypatch
) -> None:
_patch_opencloud(monkeypatch)
zip_path = tmp_path / "batch.zip"
names = [f"SKO115FZ_01_123456789_20260616_X0000{i}.xml" for i in range(1, 4)]
with zipfile.ZipFile(zip_path, "w") as archive:
for name in names:
archive.writestr(name, _valid_xml())
loose = tmp_path / "SKO115FZ_01_123456789_20260616_X00009.xml"
loose.write_text(_valid_xml(), encoding="utf-8")
calls: list[tuple[int, int, str]] = []
result = run_batch(
input_paths=[zip_path, loose],
launcher="tester",
opencloud_username="test-user",
opencloud_access_token="token-123",
progress_callback=lambda done, total, name: calls.append((done, total, name)),
)
assert result.status == "success"
assert [done for done, _, _ in calls] == [1, 2, 3, 4]
assert {total for _, total, _ in calls} == {4}
assert sorted(name for _, _, name in calls) == sorted([*names, loose.name])
def test_progress_callback_is_called_for_unreadable_file_too(
tmp_path: Path, monkeypatch
) -> None:
_patch_opencloud(monkeypatch)
loose = tmp_path / "SKO115FZ_01_123456789_20260616_X00001.xml"
loose.write_text(_valid_xml(), encoding="utf-8")
def _boom(self: Path) -> bytes:
raise OSError("disk gone")
monkeypatch.setattr(service.Path, "read_bytes", _boom)
calls: list[tuple[int, int, str]] = []
result = run_batch(
input_paths=[loose],
launcher="tester",
opencloud_username="test-user",
opencloud_access_token="token-123",
progress_callback=lambda done, total, name: calls.append((done, total, name)),
)
assert result.status == "error"
assert calls == [(1, 1, loose.name)]

View File

@ -0,0 +1,50 @@
"""Сохранение вложений Streamlit на диск не должно удваивать их в памяти.
UploadedFile — это BytesIO: Streamlit уже держит всё содержимое вложения в
памяти. getvalue() делает ещё одну полную копию каждого файла перед записью
на диск; при 230 вложениях в контейнере на 1 ГБ это лишняя пиковая нагрузка.
Копировать нужно потоково, через файловый интерфейс, не трогая getvalue().
"""
from __future__ import annotations
import io
from pathlib import Path
from main import _save_attachments
class _UploadWithoutGetValue(io.BytesIO):
"""Вложение, у которого есть только файловый интерфейс и имя."""
def __init__(self, name: str, payload: bytes) -> None:
super().__init__(payload)
self.name = name
def getvalue(self) -> bytes: # type: ignore[override]
raise AssertionError("getvalue() копирует всё вложение в память")
def test_save_attachments_streams_without_getvalue(tmp_path: Path) -> None:
uploads = [
_UploadWithoutGetValue("a.xml", b"<a/>"),
_UploadWithoutGetValue("a.xml", b"<b/>" * 1000),
]
paths = _save_attachments(uploads, tmp_path)
assert [path.name for path in paths] == ["a.xml", "a.xml"]
assert paths[0] != paths[1]
assert paths[0].read_bytes() == b"<a/>"
assert paths[1].read_bytes() == b"<b/>" * 1000
def test_save_attachments_copies_from_start_even_if_upload_was_read(
tmp_path: Path,
) -> None:
upload = _UploadWithoutGetValue("f.xml", b"0123456789")
upload.read(4)
(path,) = _save_attachments([upload], tmp_path)
assert path.read_bytes() == b"0123456789"