Fixed #30493 -- Fixed prefetch_related() for GenericRelation with different content types.

Co-Authored-By: Mariusz Felisiak <felisiak.mariusz@gmail.com>

Thanks Simon Charette for the review.
This commit is contained in:
can 2019-05-29 20:24:37 +03:00 committed by Mariusz Felisiak
parent f66021f3f7
commit dffa3e1992
2 changed files with 38 additions and 9 deletions

View File

@ -1,4 +1,6 @@
import functools import functools
import itertools
import operator
from collections import defaultdict from collections import defaultdict
from django.contrib.contenttypes.models import ContentType from django.contrib.contenttypes.models import ContentType
@ -566,19 +568,28 @@ def create_generic_related_manager(superclass, rel):
queryset._add_hints(instance=instances[0]) queryset._add_hints(instance=instances[0])
queryset = queryset.using(queryset._db or self._db) queryset = queryset.using(queryset._db or self._db)
# Group instances by content types.
query = { content_type_queries = (
'%s__pk' % self.content_type_field_name: self.content_type.id, models.Q(**{
'%s__in' % self.object_id_field_name: {obj.pk for obj in instances} '%s__pk' % self.content_type_field_name: content_type_id,
} '%s__in' % self.object_id_field_name: {obj.pk for obj in objs}
})
for content_type_id, objs in itertools.groupby(
sorted(instances, key=lambda obj: self.get_content_type(obj).pk),
lambda obj: self.get_content_type(obj).pk,
)
)
query = functools.reduce(operator.or_, content_type_queries)
# We (possibly) need to convert object IDs to the type of the # We (possibly) need to convert object IDs to the type of the
# instances' PK in order to match up instances: # instances' PK in order to match up instances:
object_id_converter = instances[0]._meta.pk.to_python object_id_converter = instances[0]._meta.pk.to_python
return ( return (
queryset.filter(**query), queryset.filter(query),
lambda relobj: object_id_converter(getattr(relobj, self.object_id_field_name)), lambda relobj: (
lambda obj: obj.pk, object_id_converter(getattr(relobj, self.object_id_field_name)),
relobj.content_type_id
),
lambda obj: (obj.pk, self.get_content_type(obj).pk),
False, False,
self.prefetch_cache_name, self.prefetch_cache_name,
False, False,

View File

@ -546,6 +546,24 @@ class GenericRelationsTests(TestCase):
platypus.tags.remove(weird_tag) platypus.tags.remove(weird_tag)
self.assertSequenceEqual(platypus.tags.all(), [furry_tag]) self.assertSequenceEqual(platypus.tags.all(), [furry_tag])
def test_prefetch_related_different_content_types(self):
TaggedItem.objects.create(content_object=self.platypus, tag='prefetch_tag_1')
TaggedItem.objects.create(
content_object=Vegetable.objects.create(name='Broccoli'),
tag='prefetch_tag_2',
)
TaggedItem.objects.create(
content_object=Animal.objects.create(common_name='Bear'),
tag='prefetch_tag_3',
)
qs = TaggedItem.objects.filter(
tag__startswith='prefetch_tag_',
).prefetch_related('content_object', 'content_object__tags')
with self.assertNumQueries(4):
tags = list(qs)
for tag in tags:
self.assertSequenceEqual(tag.content_object.tags.all(), [tag])
class ProxyRelatedModelTest(TestCase): class ProxyRelatedModelTest(TestCase):
def test_default_behavior(self): def test_default_behavior(self):