mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-07-18 18:04:53 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d863a553a | ||
|
|
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.
|
||||
|
||||
@@ -54,8 +54,12 @@ class Command(PaperlessCommand):
|
||||
else:
|
||||
documents = Document.objects.all()
|
||||
|
||||
documents = documents.only("id", "archive_filename")
|
||||
|
||||
document_ids = [
|
||||
doc.id for doc in documents if overwrite or not doc.has_archive_version
|
||||
doc.id
|
||||
for doc in documents.iterator(chunk_size=2000)
|
||||
if overwrite or not doc.has_archive_version
|
||||
]
|
||||
|
||||
try:
|
||||
|
||||
@@ -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