mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-07-18 01:44:53 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d863a553a | ||
|
|
c753091a44 |
@@ -341,6 +341,35 @@ class PaperlessCommand(RichCommand):
|
|||||||
|
|
||||||
return None
|
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(
|
def track(
|
||||||
self,
|
self,
|
||||||
iterable: Iterable[T],
|
iterable: Iterable[T],
|
||||||
@@ -367,13 +396,14 @@ class PaperlessCommand(RichCommand):
|
|||||||
for doc in self.track(documents, description="Renaming..."):
|
for doc in self.track(documents, description="Renaming..."):
|
||||||
process(doc)
|
process(doc)
|
||||||
"""
|
"""
|
||||||
|
if total is None:
|
||||||
|
total = self._get_iterable_length(iterable)
|
||||||
|
iterable = self._resolve_iterable(iterable)
|
||||||
|
|
||||||
if self.no_progress_bar:
|
if self.no_progress_bar:
|
||||||
yield from iterable
|
yield from iterable
|
||||||
return
|
return
|
||||||
|
|
||||||
if total is None:
|
|
||||||
total = self._get_iterable_length(iterable)
|
|
||||||
|
|
||||||
with self._create_progress(description) as progress:
|
with self._create_progress(description) as progress:
|
||||||
task_id = progress.add_task(description, total=total)
|
task_id = progress.add_task(description, total=total)
|
||||||
for item in iterable:
|
for item in iterable:
|
||||||
@@ -438,13 +468,14 @@ class PaperlessCommand(RichCommand):
|
|||||||
except Exception:
|
except Exception:
|
||||||
stats.failed += 1
|
stats.failed += 1
|
||||||
"""
|
"""
|
||||||
|
if total is None:
|
||||||
|
total = self._get_iterable_length(iterable)
|
||||||
|
iterable = self._resolve_iterable(iterable)
|
||||||
|
|
||||||
if self.no_progress_bar:
|
if self.no_progress_bar:
|
||||||
yield from iterable
|
yield from iterable
|
||||||
return
|
return
|
||||||
|
|
||||||
if total is None:
|
|
||||||
total = self._get_iterable_length(iterable)
|
|
||||||
|
|
||||||
stderr_console = Console(stderr=True)
|
stderr_console = Console(stderr=True)
|
||||||
|
|
||||||
# Progress is created without its own console so Live controls rendering.
|
# Progress is created without its own console so Live controls rendering.
|
||||||
|
|||||||
@@ -54,8 +54,12 @@ class Command(PaperlessCommand):
|
|||||||
else:
|
else:
|
||||||
documents = Document.objects.all()
|
documents = Document.objects.all()
|
||||||
|
|
||||||
|
documents = documents.only("id", "archive_filename")
|
||||||
|
|
||||||
document_ids = [
|
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:
|
try:
|
||||||
|
|||||||
@@ -127,12 +127,19 @@ def mock_queryset():
|
|||||||
def __init__(self, items: list):
|
def __init__(self, items: list):
|
||||||
self._items = items
|
self._items = items
|
||||||
self.count_called = False
|
self.count_called = False
|
||||||
|
self.iterator_chunk_size: int | None = None
|
||||||
|
|
||||||
def count(self) -> int:
|
def count(self) -> int:
|
||||||
self.count_called = True
|
self.count_called = True
|
||||||
return len(self._items)
|
return len(self._items)
|
||||||
|
|
||||||
def __iter__(self):
|
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)
|
return iter(self._items)
|
||||||
|
|
||||||
def __len__(self):
|
def __len__(self):
|
||||||
@@ -405,6 +412,21 @@ class TestTrack:
|
|||||||
spy.assert_called_once_with(queryset)
|
spy.assert_called_once_with(queryset)
|
||||||
assert queryset.count_called is True
|
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
|
@pytest.mark.management
|
||||||
class TestProcessParallel:
|
class TestProcessParallel:
|
||||||
|
|||||||
Reference in New Issue
Block a user