diff --git a/backend/src/chitai/services/book.py b/backend/src/chitai/services/book.py index 8cc828d..c3af57c 100644 --- a/backend/src/chitai/services/book.py +++ b/backend/src/chitai/services/book.py @@ -6,7 +6,7 @@ from collections import defaultdict import mimetypes from collections.abc import Callable -from io import BytesIO +from io import BytesIO, RawIOBase from pathlib import Path import uuid 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]): """Book service for managing book operations.""" @@ -284,7 +321,10 @@ class BookService(SQLAlchemyAsyncRepositoryService[Book]): """ 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: 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) - files = [file for book in books for file in book.files] - - 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) + stream = _ZipStream() chunk_size = 32768 # 32 KiB - while True: - chunk = buffer.read(chunk_size) - if not chunk: - break + + with zipfile.ZipFile(stream, "w", zipfile.ZIP_DEFLATED) as zip_file: + for book in books: + # 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 async def update_book( diff --git a/backend/tests/integration/test_book.py b/backend/tests/integration/test_book.py index dc74263..769615a 100644 --- a/backend/tests/integration/test_book.py +++ b/backend/tests/integration/test_book.py @@ -369,6 +369,13 @@ async def test_create_multiple_books_from_directory( 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: # raise NotImplementedError() diff --git a/backend/tests/unit/test_services/test_book_service.py b/backend/tests/unit/test_services/test_book_service.py index 92cd741..aa31a90 100644 --- a/backend/tests/unit/test_services/test_book_service.py +++ b/backend/tests/unit/test_services/test_book_service.py @@ -1,5 +1,9 @@ """Tests for BookService""" +import zipfile +from io import BytesIO +from pathlib import Path + import pytest import aiofiles.os as aios @@ -138,3 +142,63 @@ class TestBookServiceCRUD: updated_book = await books_service.get(book.id) assert len(updated_book.identifiers) == 1 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