mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-07-18 01:44:53 +00:00
Stream QuerySets in track()/track_with_stats() via .iterator()
Plain iteration over a raw QuerySet loads every row into the result cache before yielding the first item, defeating the purpose of the generator-based track() helpers. track()/track_with_stats() now auto-detect a raw QuerySet and stream it via .iterator(chunk_size=2000) while still using .count() for the progress bar total. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 5
parent
71557d7c64
commit
c753091a44
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user