From 579cf4ada728d33287eb0b7c00fe4b7af7d94322 Mon Sep 17 00:00:00 2001 From: Alex Vandiver Date: Thu, 19 Sep 2024 20:50:08 +0000 Subject: [PATCH] upload: Make local-file save_attachment_contents chunk-at-a-time. This means it does not attempt to hold large files entirely in memory when writing them to a new location on disk. --- zerver/lib/upload/local.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/zerver/lib/upload/local.py b/zerver/lib/upload/local.py index d98d77fc30..33e61f20c1 100644 --- a/zerver/lib/upload/local.py +++ b/zerver/lib/upload/local.py @@ -40,12 +40,12 @@ def write_local_file(type: Literal["avatars", "files"], path: str, file_data: by f.write(file_data) -def read_local_file(type: Literal["avatars", "files"], path: str) -> bytes: +def read_local_file(type: Literal["avatars", "files"], path: str) -> Iterator[bytes]: file_path = os.path.join(assert_is_not_none(settings.LOCAL_UPLOADS_DIR), type, path) assert_is_local_storage_path(type, file_path) with open(file_path, "rb") as f: - return f.read() + yield from iter(lambda: f.read(4 * 1024 * 1024), b"") def delete_local_file(type: Literal["avatars", "files"], path: str) -> bool: @@ -99,7 +99,8 @@ class LocalUploadBackend(ZulipUploadBackend): @override def save_attachment_contents(self, path_id: str, filehandle: IO[bytes]) -> None: - filehandle.write(read_local_file("files", path_id)) + for chunk in read_local_file("files", path_id): + filehandle.write(chunk) @override def attachment_vips_source(self, path_id: str) -> StreamingSourceWithSize: @@ -134,7 +135,7 @@ class LocalUploadBackend(ZulipUploadBackend): @override def get_avatar_contents(self, file_path: str) -> tuple[bytes, str]: - image_data = read_local_file("avatars", file_path + ".original") + image_data = b"".join(read_local_file("avatars", file_path + ".original")) content_type = guess_type(file_path)[0] return image_data, content_type or "application/octet-stream"