From fff14736f1cc594f79ea0e2656b8375d837b7aca Mon Sep 17 00:00:00 2001 From: John Parton Date: Thu, 27 Jul 2023 16:38:12 -0500 Subject: [PATCH] Fixed #34331 -- Added QuerySet.aiterator() support for prefetch_related(). --- django/db/models/query.py | 38 ++++++++++++++++++++++++------ docs/ref/models/querysets.txt | 13 +++++++--- docs/releases/5.0.txt | 6 +++++ tests/async/test_async_queryset.py | 18 ++++++++------ 4 files changed, 58 insertions(+), 17 deletions(-) diff --git a/django/db/models/query.py b/django/db/models/query.py index 395ba6e404..95323b6f9b 100644 --- a/django/db/models/query.py +++ b/django/db/models/query.py @@ -544,19 +544,36 @@ class QuerySet(AltersData): An asynchronous iterator over the results from applying this QuerySet to the database. """ - if self._prefetch_related_lookups: - raise NotSupportedError( - "Using QuerySet.aiterator() after prefetch_related() is not supported." - ) if chunk_size <= 0: raise ValueError("Chunk size must be strictly positive.") use_chunked_fetch = not connections[self.db].settings_dict.get( "DISABLE_SERVER_SIDE_CURSORS" ) - async for item in self._iterable_class( + iterable = self._iterable_class( self, chunked_fetch=use_chunked_fetch, chunk_size=chunk_size - ): - yield item + ) + if self._prefetch_related_lookups: + results = [] + + async for item in iterable: + results.append(item) + if len(results) >= chunk_size: + await aprefetch_related_objects( + results, *self._prefetch_related_lookups + ) + for result in results: + yield result + results.clear() + + if results: + await aprefetch_related_objects( + results, *self._prefetch_related_lookups + ) + for result in results: + yield result + else: + async for item in iterable: + yield item def aggregate(self, *args, **kwargs): """ @@ -2387,6 +2404,13 @@ def prefetch_related_objects(model_instances, *related_lookups): obj_list = new_obj_list +async def aprefetch_related_objects(model_instances, *related_lookups): + """See prefetch_related_objects().""" + return await sync_to_async(prefetch_related_objects)( + model_instances, *related_lookups + ) + + def get_prefetcher(instance, through_attr, to_attr): """ For the attribute 'through_attr' on the given instance, find diff --git a/docs/ref/models/querysets.txt b/docs/ref/models/querysets.txt index 92cf3baaec..8a9733786a 100644 --- a/docs/ref/models/querysets.txt +++ b/docs/ref/models/querysets.txt @@ -2579,10 +2579,10 @@ evaluated will force it to evaluate again, repeating the query. long as ``chunk_size`` is given. Larger values will necessitate fewer queries to accomplish the prefetching at the cost of greater memory usage. -.. note:: +.. versionchanged:: 5.0 - ``aiterator()`` is *not* compatible with previous calls to - ``prefetch_related()``. + Support for ``aiterator()`` with previous calls to ``prefetch_related()`` + was added. On some databases (e.g. Oracle, `SQLite `_), the maximum number @@ -4073,6 +4073,9 @@ attribute: ------------------------------ .. function:: prefetch_related_objects(model_instances, *related_lookups) +.. function:: aprefetch_related_objects(model_instances, *related_lookups) + +*Asynchronous version*: ``aprefetch_related_objects()`` Prefetches the given lookups on an iterable of model instances. This is useful in code that receives a list of model instances as opposed to a ``QuerySet``; @@ -4091,6 +4094,10 @@ When using multiple databases with ``prefetch_related_objects``, the prefetch query will use the database associated with the model instance. This can be overridden by using a custom queryset in a related lookup. +.. versionchanged:: 5.0 + + ``aprefetch_related_objects()`` function was added. + ``FilteredRelation()`` objects ------------------------------ diff --git a/docs/releases/5.0.txt b/docs/releases/5.0.txt index e7fcc6e07f..cceba667f9 100644 --- a/docs/releases/5.0.txt +++ b/docs/releases/5.0.txt @@ -368,6 +368,12 @@ Models :func:`~django.shortcuts.aget_list_or_404` asynchronous shortcuts allow asynchronous getting objects. +* The new :func:`~django.db.models.aprefetch_related_objects` function allows + asynchronous prefetching of model instances. + +* :meth:`.QuerySet.aiterator` now supports previous calls to + ``prefetch_related()``. + Pagination ~~~~~~~~~~ diff --git a/tests/async/test_async_queryset.py b/tests/async/test_async_queryset.py index e83157f5bd..374b4576f9 100644 --- a/tests/async/test_async_queryset.py +++ b/tests/async/test_async_queryset.py @@ -5,10 +5,10 @@ from datetime import datetime from asgiref.sync import async_to_sync, sync_to_async from django.db import NotSupportedError, connection -from django.db.models import Sum +from django.db.models import Prefetch, Sum from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature -from .models import SimpleModel +from .models import RelatedModel, SimpleModel class AsyncQuerySetTest(TestCase): @@ -26,6 +26,9 @@ class AsyncQuerySetTest(TestCase): field=3, created=datetime(2022, 1, 1, 0, 0, 2), ) + cls.r1 = RelatedModel.objects.create(simple=cls.s1) + cls.r2 = RelatedModel.objects.create(simple=cls.s2) + cls.r3 = RelatedModel.objects.create(simple=cls.s3) @staticmethod def _get_db_feature(connection_, feature_name): @@ -48,11 +51,12 @@ class AsyncQuerySetTest(TestCase): self.assertCountEqual(results, [self.s1, self.s2, self.s3]) async def test_aiterator_prefetch_related(self): - qs = SimpleModel.objects.prefetch_related("relatedmodels").aiterator() - msg = "Using QuerySet.aiterator() after prefetch_related() is not supported." - with self.assertRaisesMessage(NotSupportedError, msg): - async for m in qs: - pass + results = [] + async for s in SimpleModel.objects.prefetch_related( + Prefetch("relatedmodel_set", to_attr="prefetched_relatedmodel") + ).aiterator(): + results.append(s.prefetched_relatedmodel) + self.assertCountEqual(results, [[self.r1], [self.r2], [self.r3]]) async def test_aiterator_invalid_chunk_size(self): msg = "Chunk size must be strictly positive."