Refactor: extract QuerySetStream, shared by _StreamedDocuments and _DocumentViewerStream

src/documents/search/_backend.py

@@ -39,6 +39,7 @@ from documents.search._tokenizer import ascii_fold
 from documents.search._tokenizer import autocomplete_tokens
 from documents.search._tokenizer import register_tokenizers
 from documents.utils import IterWrapper
+from documents.utils import QuerySetStream
 from documents.utils import identity
 
 if TYPE_CHECKING:
@@ -1035,14 +1036,15 @@ _EMPTY_VIEWER_GRANT: Final[ViewerGrant] = ViewerGrant(
 )
 
 
-class _DocumentViewerStream:
+class _DocumentViewerStream(QuerySetStream["Document"]):
     """Yield document permission data while batch-loading grants.
 
     Viewer permissions are fetched in batches (see
     ``_bulk_get_viewer_permissions``), but documents are yielded individually so a
     progress bar wrapped around this stream advances per document rather than
-    jumping a whole chunk at a time. ``__len__`` lets the progress helper still
-    discover the total (it inspects ``QuerySet``/``Sized``).
+    jumping a whole chunk at a time. ``__len__`` (inherited from
+    ``QuerySetStream``) lets the progress helper still discover the total (it
+    inspects ``QuerySet``/``Sized``).
 
     The viewer and group ids travel with each document in the yielded pair
     rather than through a separate mutable attribute, so the pairing survives
@@ -1051,18 +1053,11 @@ class _DocumentViewerStream:
     generator in lock-step.
     """
 
-    def __init__(self, documents: QuerySet[Document], *, chunk_size: int) -> None:
-        self._documents = documents
-        self._chunk_size = chunk_size
-
-    def __len__(self) -> int:
-        return self._documents.count()
-
     def __iter__(self) -> Iterator[tuple[Document, ViewerGrant]]:
         # iterator(chunk_size=…) streams from a server-side cursor instead of
         # materialising the whole queryset in memory; since Django 4.1 it still
         # honours prefetch_related, running the prefetches one batch at a time.
-        documents = self._documents.iterator(chunk_size=self._chunk_size)
+        documents = self._queryset.iterator(chunk_size=self._chunk_size)
         for chunk in chunked(documents, self._chunk_size):
             grants_by_pk = _bulk_get_viewer_permissions([doc.pk for doc in chunk])
             for doc in chunk:

src/documents/tests/test_utils.py

@@ -0,0 +1,35 @@
+import pytest_mock
+
+from documents.utils import QuerySetStream
+
+
+class TestQuerySetStream:
+    def test_len_and_iter_delegate_to_streaming_queryset_methods(
+        self,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - A mock queryset
+        WHEN:
+            - A QuerySetStream wrapping it is measured and iterated
+        THEN:
+            - len() uses count() (not a materializing len()), and iteration
+              uses .iterator(chunk_size=...) (not plain iteration, which
+              would materialize the whole queryset, plus any prefetch
+              caches, into Django's own result cache at once)
+        """
+        mock_queryset = mocker.MagicMock()
+        mock_queryset.count.return_value = 42
+        mock_queryset.iterator.return_value = iter(["row-1", "row-2"])
+        streamed = QuerySetStream(mock_queryset, chunk_size=1000)
+
+        assert len(streamed) == 42
+        assert list(streamed) == ["row-1", "row-2"]
+        # count.call_count isn't asserted exactly: list()'s own size-hint
+        # optimization calls len(streamed) again internally, on top of the
+        # explicit len() call above -- both legitimately delegate to
+        # count(), so only the delegation itself (not the call count) is
+        # the thing being verified here.
+        mock_queryset.count.assert_called_with()
+        mock_queryset.iterator.assert_called_once_with(chunk_size=1000)

src/documents/utils.py

@@ -3,16 +3,24 @@ import logging
 import shutil
 from collections.abc import Callable
 from collections.abc import Iterable
+from collections.abc import Iterator
 from os import utime
 from pathlib import Path
 from subprocess import CompletedProcess
 from subprocess import run
+from typing import TYPE_CHECKING
+from typing import Generic
 from typing import TypeVar
 
 from django.conf import settings
 from PIL import Image
 
+if TYPE_CHECKING:
+    from django.db.models import Model
+    from django.db.models import QuerySet
+
 _T = TypeVar("_T")
+_M = TypeVar("_M", bound="Model")
 
 # A function that wraps an iterable — typically used to inject a progress bar.
 IterWrapper = Callable[[Iterable[_T]], Iterable[_T]]
@@ -23,6 +31,40 @@ def identity(iterable: Iterable[_T]) -> Iterable[_T]:
     return iterable
 
 
+class QuerySetStream(Generic[_M]):
+    """Stream a QuerySet via .iterator(chunk_size=...) instead of
+    materializing it (plus any prefetch caches) all at once, while still
+    supporting len() via count() so a progress bar wrapped around this
+    (e.g. via IterWrapper) shows a real total instead of falling back to
+    indeterminate.
+
+    Plain QuerySet iteration (``for row in queryset:``) is not lazy: Django
+    fetches every matching row in one query and caches the fully-hydrated
+    result in the queryset's own ``_result_cache`` before yielding the
+    first item -- wrapping that in a progress bar or any other iterable
+    adapter doesn't change this, since none of them alter how the
+    underlying queryset produces items. ``.iterator(chunk_size=...)`` is
+    the specific Django API that bypasses ``_result_cache`` and streams
+    from a server-side cursor instead, discarding each chunk once consumed
+    (and, since Django 4.1, still honours ``prefetch_related``, running the
+    prefetches one batch at a time rather than for the whole queryset).
+
+    Subclass to layer additional per-batch work on top (see
+    ``documents.search._backend._DocumentViewerStream``) by overriding
+    ``__iter__`` -- ``__len__`` and the constructor are inherited for free.
+    """
+
+    def __init__(self, queryset: "QuerySet[_M]", *, chunk_size: int) -> None:
+        self._queryset = queryset
+        self._chunk_size = chunk_size
+
+    def __len__(self) -> int:
+        return self._queryset.count()
+
+    def __iter__(self) -> Iterator[_M]:
+        return iter(self._queryset.iterator(chunk_size=self._chunk_size))
+
+
 def _coerce_to_path(
     source: Path | str,
     dest: Path | str,

src/paperless_ai/indexing.py

@@ -1,6 +1,5 @@
 import logging
 from collections.abc import Iterable
-from collections.abc import Iterator
 from contextlib import contextmanager
 from datetime import timedelta
 from typing import TYPE_CHECKING
@@ -15,6 +14,7 @@ from filelock import Timeout
 from documents.models import Document
 from documents.models import PaperlessTask
 from documents.utils import IterWrapper
+from documents.utils import QuerySetStream
 from documents.utils import identity
 from paperless.config import AIConfig
 from paperless_ai.db import db_connection_released
@@ -23,7 +23,6 @@ from paperless_ai.embedding import get_configured_model_name
 from paperless_ai.embedding import get_embedding_model
 
 if TYPE_CHECKING:
-    from django.db.models import QuerySet
     from llama_index.core.schema import BaseNode
 
     from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
@@ -35,35 +34,11 @@ RAG_NUM_OUTPUT = 512
 RAG_CHUNK_OVERLAP = 200
 
 # update_llm_index(): row count per .iterator() batch when streaming
-# documents for a rebuild/update, matching _DocumentViewerStream's chunk
-# size in documents/search/_backend.py.
+# documents for a rebuild/update via QuerySetStream, matching
+# _DocumentViewerStream's chunk size in documents/search/_backend.py.
 _INDEX_STREAM_CHUNK_SIZE = 1000
 
 
-class _StreamedDocuments:
-    """A thin QuerySet wrapper that streams via ``.iterator()`` instead of
-    materializing every row (plus its ``content`` and prefetch caches) into
-    memory at once, while still supporting ``len()`` so ``iter_wrapper``'s
-    progress bar shows a real total instead of falling back to indeterminate.
-    Same shape as ``documents/search/_backend.py``'s ``_DocumentViewerStream``,
-    just without that class's extra per-batch permission lookup -- nothing
-    here needs one.
-    """
-
-    def __init__(self, documents: "QuerySet[Document]") -> None:
-        self._documents = documents
-
-    def __len__(self) -> int:
-        return self._documents.count()
-
-    def __iter__(self) -> Iterator[Document]:
-        # iterator(chunk_size=...) streams from a server-side cursor instead
-        # of materializing the whole queryset in memory; since Django 4.1 it
-        # still honours prefetch_related, running the prefetches one batch
-        # at a time.
-        return iter(self._documents.iterator(chunk_size=_INDEX_STREAM_CHUNK_SIZE))
-
-
 def queue_llm_index_update_if_needed(*, rebuild: bool, reason: str) -> bool:
     # NOTE: The check-then-enqueue sequence below is non-atomic (TOCTOU): two
     # concurrent workers can both observe no running task and both enqueue a
@@ -416,7 +391,9 @@ def update_llm_index(
         if rebuild or not store.table_exists():
             logger.info("Rebuilding LLM index.")
             store.drop_table()
-            for document in iter_wrapper(_StreamedDocuments(documents)):
+            for document in iter_wrapper(
+                QuerySetStream(documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE),
+            ):
                 nodes = build_document_node(document, chunk_size=chunk_size)
                 _embed_nodes(nodes, embed_model)
                 store.add(nodes)
@@ -429,7 +406,9 @@ def update_llm_index(
             )
             existing = store.get_modified_times()
             changed = 0
-            for document in iter_wrapper(_StreamedDocuments(scoped_documents)):
+            for document in iter_wrapper(
+                QuerySetStream(scoped_documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE),
+            ):
                 doc_id = str(document.id)
                 if existing.get(doc_id) == document.modified.isoformat():
                     continue

src/paperless_ai/tests/test_ai_indexing.py

@@ -186,45 +186,6 @@ def test_truncate_embedding_query_returns_single_chunk() -> None:
     assert "word199" not in result
 
 
-class TestStreamedDocuments:
-    """_StreamedDocuments streams via .iterator() instead of materializing
-    the whole queryset (plus its content and prefetch caches) in memory at
-    once, while still supporting len() so a progress bar wrapped around it
-    shows a real total.
-    """
-
-    def test_len_and_iter_delegate_to_streaming_queryset_methods(
-        self,
-        mocker: pytest_mock.MockerFixture,
-    ) -> None:
-        """
-        GIVEN:
-            - A mock queryset
-        WHEN:
-            - A _StreamedDocuments wrapping it is measured and iterated
-        THEN:
-            - len() uses count() (not a materializing len()), and iteration
-              uses .iterator(chunk_size=...) (not plain iteration, which
-              would materialize prefetches for the whole queryset at once)
-        """
-        mock_queryset = mocker.MagicMock()
-        mock_queryset.count.return_value = 42
-        mock_queryset.iterator.return_value = iter(["doc-1", "doc-2"])
-        streamed = indexing._StreamedDocuments(mock_queryset)
-
-        assert len(streamed) == 42
-        assert list(streamed) == ["doc-1", "doc-2"]
-        # count.call_count isn't asserted exactly: list()'s own size-hint
-        # optimization calls len(streamed) again internally, on top of the
-        # explicit len() call above -- both legitimately delegate to
-        # count(), so only the delegation itself (not the call count) is
-        # the thing being verified here.
-        mock_queryset.count.assert_called_with()
-        mock_queryset.iterator.assert_called_once_with(
-            chunk_size=indexing._INDEX_STREAM_CHUNK_SIZE,
-        )
-
-
 @pytest.mark.django_db
 def test_update_llm_index(
     temp_llm_index_dir: Path,