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:
@@ -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(
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user