Commit 69f7db15 authored by Tai Lee's avatar Tai Lee Committed by Anssi Kääriäinen
Browse files

Fixed #16436 -- defer + annotate + select_related crash 

Correctly calculate the ``aggregate_start`` offset from loaded fields,
if any are deferred, instead of ``self.query.select`` which includes all
fields on the model.

Also made some PEP 8 fixes.
parent bb80d297

django/db/models/sql/compiler.py

+2 −1
Original line number Diff line number Diff line
@@ -729,7 +729,8 @@ class SQLCompiler(object): 
                    row = self.resolve_columns(row, fields)



                if has_aggregate_select:

                    aggregate_start = len(self.query.extra_select) + len(self.query.select)

                    loaded_fields = self.query.get_loaded_field_names().get(self.query.model, set()) or self.query.select

                    aggregate_start = len(self.query.extra_select) + len(loaded_fields)

                    aggregate_end = aggregate_start + len(self.query.aggregate_select)

                    row = tuple(row[:aggregate_start]) + tuple([

                        self.query.resolve_aggregate(value, aggregate, self.connection)

tests/defer_regress/models.py

+19 −0
Original line number Diff line number Diff line
@@ -62,3 +62,22 @@ class OneToOneItem(models.Model): 
class ItemAndSimpleItem(models.Model):

    item = models.ForeignKey(Item)

    simple = models.ForeignKey(SimpleItem)



class Profile(models.Model):

    profile1 = models.TextField(default='profile1')



class Location(models.Model):

    location1 = models.TextField(default='location1')



class Item(models.Model):

    pass



class Request(models.Model):

    profile = models.ForeignKey(Profile, null=True, blank=True)

    location = models.ForeignKey(Location)

    items = models.ManyToManyField(Item)



    request1 = models.TextField(default='request1')

    request2 = models.TextField(default='request2')

    request3 = models.TextField(default='request3')

    request4 = models.TextField(default='request4')

tests/defer_regress/tests.py

+34 −7
Original line number Diff line number Diff line
@@ -8,8 +8,9 @@ from django.db.models import Count 
from django.db.models.loading import cache

from django.test import TestCase



from .models import (ResolveThis, Item, RelatedItem, Child, Leaf, Proxy,

    SimpleItem, Feature, ItemAndSimpleItem, OneToOneItem, SpecialFeature)

from .models import (

    ResolveThis, Item, RelatedItem, Child, Leaf, Proxy, SimpleItem, Feature,

    ItemAndSimpleItem, OneToOneItem, SpecialFeature, Location, Request)





class DeferRegressionTest(TestCase):
@@ -76,7 +77,9 @@ class DeferRegressionTest(TestCase): 
        self.assertEqual(results[0].child.name, "c1")

        self.assertEqual(results[0].second_child.name, "c2")



        results = Leaf.objects.only("name", "child", "second_child", "child__name", "second_child__name").select_related()

        results = Leaf.objects.only(

            "name", "child", "second_child", "child__name", "second_child__name"

        ).select_related()

        self.assertEqual(results[0].child.name, "c1")

        self.assertEqual(results[0].second_child.name, "c2")
@@ -144,6 +147,14 @@ class DeferRegressionTest(TestCase): 
                cache.get_app("defer_regress"), include_deferred=True))

        )



        # Regression for #16409 - make sure defer() and only() work with annotate()

        self.assertIsInstance(

            list(SimpleItem.objects.annotate(Count('feature')).defer('name')),

            list)

        self.assertIsInstance(

            list(SimpleItem.objects.annotate(Count('feature')).only('name')),

            list)



    def test_only_and_defer_usage_on_proxy_models(self):

        # Regression for #15790 - only() broken for proxy models

        proxy = Proxy.objects.create(name="proxy", value=42)
@@ -160,7 +171,7 @@ class DeferRegressionTest(TestCase): 
        self.assertEqual(dp.value, proxy.value, msg=msg)



    def test_resolve_columns(self):

        rt = ResolveThis.objects.create(num=5.0, name='Foobar')

        ResolveThis.objects.create(num=5.0, name='Foobar')

        qs = ResolveThis.objects.defer('num')

        self.assertEqual(1, qs.count())

        self.assertEqual('Foobar', qs[0].name)
@@ -196,7 +207,7 @@ class DeferRegressionTest(TestCase): 
        item1 = Item.objects.create(name="first", value=47)

        item2 = Item.objects.create(name="second", value=42)

        simple = SimpleItem.objects.create(name="simple", value="23")

        related = ItemAndSimpleItem.objects.create(item=item1, simple=simple)

        ItemAndSimpleItem.objects.create(item=item1, simple=simple)



        obj = ItemAndSimpleItem.objects.defer('item').select_related('simple').get()

        self.assertEqual(obj.item, item1)
@@ -223,7 +234,23 @@ class DeferRegressionTest(TestCase): 


    def test_deferred_class_factory(self):

        from django.db.models.query_utils import deferred_class_factory

        new_class = deferred_class_factory(Item,

        new_class = deferred_class_factory(

            Item,

            ('this_is_some_very_long_attribute_name_so_modelname_truncation_is_triggered',))

        self.assertEqual(new_class.__name__,

        self.assertEqual(

            new_class.__name__,

            'Item_Deferred_this_is_some_very_long_attribute_nac34b1f495507dad6b02e2cb235c875e')



class DeferAnnotateSelectRelatedTest(TestCase):

    def test_defer_annotate_select_related(self):

        location = Location.objects.create()

        Request.objects.create(location=location)

        self.assertIsInstance(list(Request.objects

            .annotate(Count('items')).select_related('profile', 'location')

            .only('profile', 'location')), list)

        self.assertIsInstance(list(Request.objects

            .annotate(Count('items')).select_related('profile', 'location')

            .only('profile__profile1', 'location__location1')), list)

        self.assertIsInstance(list(Request.objects

            .annotate(Count('items')).select_related('profile', 'location')

            .defer('request1', 'request2', 'request3', 'request4')), list)