diff --git a/src/documents/management/commands/base.py b/src/documents/management/commands/base.py index 1d76460c0..006b933f4 100644 --- a/src/documents/management/commands/base.py +++ b/src/documents/management/commands/base.py @@ -341,6 +341,35 @@ class PaperlessCommand(RichCommand): return None + def _resolve_iterable( + self, + iterable: Iterable[T], + *, + chunk_size: int = 2000, + ) -> Iterable[T]: + """ + Convert a raw Django QuerySet into a server-side iterator. + + Plain iteration over a QuerySet (`for x in queryset`) loads every + row into the queryset's result cache before yielding the first + item, defeating the purpose of streaming via track()/ + track_with_stats(). Wrapping it in .iterator() streams rows from + the database instead. Non-QuerySet iterables (generators, lists, + already-wrapped iterators) are returned unchanged. + + Args: + iterable: The items to iterate over. + chunk_size: Row batch size passed to QuerySet.iterator(). + + Returns: + A QuerySet.iterator() if given a QuerySet, else the original + iterable. + """ + if isinstance(iterable, QuerySet): + return iterable.iterator(chunk_size=chunk_size) + + return iterable + def track( self, iterable: Iterable[T], @@ -367,13 +396,14 @@ class PaperlessCommand(RichCommand): for doc in self.track(documents, description="Renaming..."): process(doc) """ + if total is None: + total = self._get_iterable_length(iterable) + iterable = self._resolve_iterable(iterable) + if self.no_progress_bar: yield from iterable return - if total is None: - total = self._get_iterable_length(iterable) - with self._create_progress(description) as progress: task_id = progress.add_task(description, total=total) for item in iterable: @@ -438,13 +468,14 @@ class PaperlessCommand(RichCommand): except Exception: stats.failed += 1 """ + if total is None: + total = self._get_iterable_length(iterable) + iterable = self._resolve_iterable(iterable) + if self.no_progress_bar: yield from iterable return - if total is None: - total = self._get_iterable_length(iterable) - stderr_console = Console(stderr=True) # Progress is created without its own console so Live controls rendering. diff --git a/src/documents/tests/management/test_management_base_cmd.py b/src/documents/tests/management/test_management_base_cmd.py index 572bcd52d..cb3340b33 100644 --- a/src/documents/tests/management/test_management_base_cmd.py +++ b/src/documents/tests/management/test_management_base_cmd.py @@ -127,12 +127,19 @@ def mock_queryset(): def __init__(self, items: list): self._items = items self.count_called = False + self.iterator_chunk_size: int | None = None def count(self) -> int: self.count_called = True return len(self._items) def __iter__(self): + raise AssertionError( + "querysets should be streamed via .iterator(), not iterated directly", + ) + + def iterator(self, chunk_size: int | None = None): + self.iterator_chunk_size = chunk_size return iter(self._items) def __len__(self): @@ -405,6 +412,21 @@ class TestTrack: spy.assert_called_once_with(queryset) assert queryset.count_called is True + def test_streams_querysets_via_iterator( + self, + simple_command: SimpleCommand, + mock_queryset, + ) -> None: + """Verify track() streams querysets via .iterator() rather than + materializing the full result cache with a plain for-loop.""" + simple_command.no_progress_bar = False + queryset = mock_queryset([1, 2, 3]) + + result = list(simple_command.track(queryset)) + + assert result == [1, 2, 3] + assert queryset.iterator_chunk_size == 2000 + @pytest.mark.management class TestProcessParallel: