Fixed #34331 -- Added QuerySet.aiterator() support for prefetch_related().

This commit is contained in:
John Parton 2023-07-27 16:38:12 -05:00 committed by Mariusz Felisiak
parent 1ad7761ee6
commit fff14736f1
4 changed files with 58 additions and 17 deletions

View file

@ -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."