fix: resolve file paths and stream the multi-book zip download

file.path is relative to book.path, so the archive matched nothing on disk and
came out empty. Also namespace entries per book to stop filename collisions, and
stream the zip instead of buffering it whole.
This commit is contained in:
2026-08-12 15:17:48 -04:00
parent 540522e828
commit bd8d68b9ba
3 changed files with 146 additions and 16 deletions
+75 -16
View File
@@ -6,7 +6,7 @@ from collections import defaultdict
import mimetypes import mimetypes
from collections.abc import Callable from collections.abc import Callable
from io import BytesIO from io import BytesIO, RawIOBase
from pathlib import Path from pathlib import Path
import uuid import uuid
import zipfile import zipfile
@@ -57,6 +57,43 @@ from chitai.services.utils import (
) )
class _ZipStream(RawIOBase):
"""
A write-only sink that hands whatever `ZipFile` writes back to the caller.
`ZipFile` wants a file object, but a download handler wants chunks it can yield.
This collects the bytes `ZipFile` produces so the generator driving it can drain
them as they appear, instead of building the whole archive first. Reporting
itself as unseekable makes `ZipFile` emit data descriptors rather than seeking
back to patch entry headers.
"""
def __init__(self) -> None:
self._chunks: list[bytes] = []
self._position = 0
def writable(self) -> bool:
return True
def seekable(self) -> bool:
return False
def tell(self) -> int:
return self._position
def write(self, data: Any) -> int:
chunk = bytes(data)
self._chunks.append(chunk)
self._position += len(chunk)
return len(chunk)
def drain(self) -> bytes:
"""Take everything written since the last drain."""
data = b"".join(self._chunks)
self._chunks.clear()
return data
class BookService(SQLAlchemyAsyncRepositoryService[Book]): class BookService(SQLAlchemyAsyncRepositoryService[Book]):
"""Book service for managing book operations.""" """Book service for managing book operations."""
@@ -284,7 +321,10 @@ class BookService(SQLAlchemyAsyncRepositoryService[Book]):
""" """
Get all selected book files as a compressed zip file. Get all selected book files as a compressed zip file.
Streams the zip file in chunks to avoid loading the entire file into memory. Streams the zip file in chunks to avoid loading the entire archive into memory.
Each book gets its own directory in the archive, so two books that happen to
share a filename do not overwrite one another. Files that have a row but no
longer exist on disk are skipped.
Args: Args:
book_ids: List of book IDs to include in the zip. book_ids: List of book IDs to include in the zip.
@@ -295,21 +335,40 @@ class BookService(SQLAlchemyAsyncRepositoryService[Book]):
""" """
books = await self.list(Book.id.in_(book_ids), Book.library_id == library_id) books = await self.list(Book.id.in_(book_ids), Book.library_id == library_id)
files = [file for book in books for file in book.files] stream = _ZipStream()
buffer = BytesIO()
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
for file in files:
path = Path(file.path)
if path.exists():
zip_file.write(path, arcname=path.name)
buffer.seek(0)
chunk_size = 32768 # 32 KiB chunk_size = 32768 # 32 KiB
while True:
chunk = buffer.read(chunk_size) with zipfile.ZipFile(stream, "w", zipfile.ZIP_DEFLATED) as zip_file:
if not chunk: for book in books:
break # Books without files have no path, and `file.path` is stored relative
# to it, so there is nothing to resolve against.
if book.path is None:
continue
book_path = Path(book.path)
for file in book.files:
path = book_path / file.path
if not await aios.path.isfile(path):
continue
# Namespaced by the book's own directory to avoid collisions.
info = zipfile.ZipInfo.from_file(
path, arcname=f"{book_path.name}/{Path(file.path).name}"
)
info.compress_type = zipfile.ZIP_DEFLATED
with zip_file.open(info, "w") as entry:
async with aiofiles.open(path, "rb") as source:
while content := await source.read(chunk_size):
entry.write(content)
if chunk := stream.drain():
yield chunk
if chunk := stream.drain():
yield chunk
# The central directory, written when ZipFile closes.
if chunk := stream.drain():
yield chunk yield chunk
async def update_book( async def update_book(
+7
View File
@@ -369,6 +369,13 @@ async def test_create_multiple_books_from_directory(
assert len(data.get("items") or data.get("data")) >= 1 assert len(data.get("items") or data.get("data")) >= 1
# NOTE: the multi-book ZIP download is covered at the service level, in
# tests/unit/test_services/test_book_service.py. Driving `/books/download` through
# AsyncTestClient hangs in fixture teardown: it is the only `Stream` endpoint in the
# app, and the test transport never sends the `http.disconnect` that Litestar's
# streaming response waits on, so the app's lifespan shutdown never completes.
# async def test_delete_book_metadata(authenticated_client: AsyncClient) -> None: # async def test_delete_book_metadata(authenticated_client: AsyncClient) -> None:
# raise NotImplementedError() # raise NotImplementedError()
@@ -1,5 +1,9 @@
"""Tests for BookService""" """Tests for BookService"""
import zipfile
from io import BytesIO
from pathlib import Path
import pytest import pytest
import aiofiles.os as aios import aiofiles.os as aios
@@ -138,3 +142,63 @@ class TestBookServiceCRUD:
updated_book = await books_service.get(book.id) updated_book = await books_service.get(book.id)
assert len(updated_book.identifiers) == 1 assert len(updated_book.identifiers) == 1
assert updated_book.identifiers[0].value == "9780261102999" assert updated_book.identifiers[0].value == "9780261102999"
async def test_get_files_zips_every_book_file(
self, books_service: BookService, test_library: m.Library
) -> None:
"""The multi-book download must contain the real bytes of each file.
`file.path` holds a bare filename relative to `book.path`, so resolving it on
its own matched nothing on disk and yielded a valid but empty archive.
"""
contents = {
"Dracula.epub": b"epub-payload-" * 512,
"Dracula.pdf": b"pdf-payload-" * 512,
}
book_data = BookCreate(
library_id=test_library.id,
title="Dracula",
authors=["Bram Stoker"],
pages=418,
)
book = await books_service.to_model_on_create(book_data.model_dump())
assert isinstance(book, m.Book)
book.path = f"{test_library.root_path}/Bram Stoker/Dracula"
await aios.makedirs(book.path) # type: ignore[arg-type]
for name, payload in contents.items():
Path(book.path, name).write_bytes(payload)
book.files.append(
m.FileMetadata(
path=name,
size=len(payload),
hash=f"hash-{name}",
content_type=None,
)
)
books_service.repository.session.add(book)
await books_service.repository.session.commit()
archive_bytes = b"".join(
[
chunk
async for chunk in books_service.get_files([book.id], test_library.id)
]
)
with zipfile.ZipFile(BytesIO(archive_bytes)) as archive:
assert archive.testzip() is None
names = archive.namelist()
assert len(names) == len(contents)
# Entries are namespaced by the book's directory, so two books sharing a
# filename cannot overwrite one another.
assert {Path(name).parent.name for name in names} == {"Dracula"}
for name, payload in contents.items():
entry = next(n for n in names if Path(n).name == name)
assert archive.read(entry) == payload