Ticket #18177: prefetch_related_cache_with_tests.patch

tests/modeltests/existing_related_instances/tests.py

@@ +1 @@
+from __future__ import absolute_import
+
+from django.test import TestCase
+
+from .models import Tournament, Pool, PoolStyle
+
+class EriTests(TestCase):
+    fixtures = ['tournament.json',]
+
+    def test_fk(self):
+        with self.assertNumQueries(2):
+            tournament = Tournament.objects.get(pk=1)
+            pool = tournament.pool_set.all()[0]
+            self.assertIs(tournament, pool.tournament)
+
+    def test_fk_prefetch_related(self):
+        with self.assertNumQueries(2):
+            tournament = (
+                Tournament.objects.prefetch_related('pool_set')
+                .get(pk=1)
+                )
+            pool = tournament.pool_set.all()[0]
+            self.assertIs(tournament, pool.tournament)
+
+    def test_fk_multiple_prefetch(self):
+        with self.assertNumQueries(2):
+            tournaments = list(
+                Tournament.objects.prefetch_related('pool_set')
+                )
+            pool1 = tournaments[0].pool_set.all()[0]
+            self.assertIs(tournaments[0], pool1.tournament)
+            pool2 = tournaments[1].pool_set.all()[0]
+            self.assertIs(tournaments[1], pool2.tournament)
+
+    def test_1t1(self):
+        with self.assertNumQueries(2):
+            style = PoolStyle.objects.get(pk=1)
+            pool = style.pool
+            self.assertIs(style, pool.poolstyle)
+
+    def test_1t1_select_related(self):
+        with self.assertNumQueries(1):
+            style = PoolStyle.objects.select_related('pool').get(pk=1)
+            pool = style.pool
+            self.assertIs(style, pool.poolstyle)
+
+    def test_1t1_prefetch_related(self):
+        with self.assertNumQueries(2):
+            style = PoolStyle.objects.prefetch_related('pool').get(pk=1)
+            pool = style.pool
+            self.assertIs(style, pool.poolstyle)
+
+    def test_1t1r(self):
+        with self.assertNumQueries(2):
+            pool = Pool.objects.get(pk=2)
+            style = pool.poolstyle
+            self.assertIs(pool, style.pool)
+
+    def test_1t1r_select_related(self):
+        with self.assertNumQueries(1):
+            pool = Pool.objects.select_related('poolstyle').get(pk=2)
+            style = pool.poolstyle
+            self.assertIs(pool, style.pool)
+
+    def test_1t1r_prefetch_related(self):
+        with self.assertNumQueries(2):
+            pool = Pool.objects.prefetch_related('poolstyle').get(pk=2)
+            style = pool.poolstyle
+            self.assertIs(pool, style.pool)
+
+    def test_1t1_multi_related(self):
+        with self.assertNumQueries(1):
+            poolstyles = list(PoolStyle.objects.select_related('pool'))
+            self.assertIs(poolstyles[0], poolstyles[0].pool.poolstyle)
+            self.assertIs(poolstyles[1], poolstyles[1].pool.poolstyle)
+
+    def test_1t1_multi_prefetch(self):
+        with self.assertNumQueries(2):
+            poolstyles = list(PoolStyle.objects.prefetch_related('pool'))
+            self.assertIs(poolstyles[0], poolstyles[0].pool.poolstyle)
+            self.assertIs(poolstyles[1], poolstyles[1].pool.poolstyle)
+
+    def test_1t1r_multi_related(self):
+        with self.assertNumQueries(1):
+            pools = list(Pool.objects.select_related('poolstyle'))
+            self.assertIs(pools[1], pools[1].poolstyle.pool)
+            self.assertIs(pools[2], pools[2].poolstyle.pool)
+
+    def test_1t1r_multi_prefetch(self):
+        with self.assertNumQueries(2):
+            pools = list(Pool.objects.prefetch_related('poolstyle'))
+            self.assertIs(pools[1], pools[1].poolstyle.pool)
+            self.assertIs(pools[2], pools[2].poolstyle.pool)

tests/modeltests/existing_related_instances/fixtures/tournament.json

@@ +1 @@
+[
+    {
+        "pk": 1,
+        "model": "existing_related_instances.tournament",
+        "fields": {
+            "name": "Tourney 1"
+            }
+        },
+    {
+        "pk": 2,
+        "model": "existing_related_instances.tournament",
+        "fields": {
+            "name": "Tourney 2"
+            }
+        },
+    {
+        "pk": 1,
+        "model": "existing_related_instances.pool",
+        "fields": {
+            "tournament": 1,
+            "name": "T1 Pool 1"
+            }
+        },
+    {
+        "pk": 2,
+        "model": "existing_related_instances.pool",
+        "fields": {
+            "tournament": 1,
+            "name": "T1 Pool 2"
+            }
+        },
+    {
+        "pk": 3,
+        "model": "existing_related_instances.pool",
+        "fields": {
+            "tournament": 2,
+            "name": "T2 Pool 1"
+            }
+        },
+    {
+        "pk": 4,
+        "model": "existing_related_instances.pool",
+        "fields": {
+            "tournament": 2,
+            "name": "T2 Pool 2"
+            }
+        },
+    {
+        "pk": 1,
+        "model": "existing_related_instances.poolstyle",
+        "fields": {
+            "name": "T1 Pool 2 Style",
+            "pool": 2
+            }
+        },
+    {
+        "pk": 2,
+        "model": "existing_related_instances.poolstyle",
+        "fields": {
+            "name": "T2 Pool 1 Style",
+            "pool": 3
+            }
+        }
+]
+

tests/modeltests/existing_related_instances/models.py

@@ +1 @@
+"""
+1. Existing related object instance caching
+
+These test that queries are not redone when going back through already
+explored relations.
+"""
+
+from django.db import models
+
+class Tournament(models.Model):
+    name = models.CharField(max_length=30)
+
+class Pool(models.Model):
+    name = models.CharField(max_length=30)
+    tournament = models.ForeignKey(Tournament)
+
+class PoolStyle(models.Model):
+    name = models.CharField(max_length=30)
+    pool = models.OneToOneField(Pool)
+

django/db/models/fields/related.py

@@ -84 +84 @@
     for cls, field, operation in pending_lookups.pop(key, []):
         operation(field, sender, cls)
 
+def set_eri(qs, field_name, rel_field_name, instances):
+    qs._existing_related_instances = {
+        field_name: dict(
+            (getattr(instance, rel_field_name), instance)
+            for instance in instances),
+        }
+
 signals.class_prepared.connect(do_pending_lookups)
 
 #HACK
@@ -232 +239 @@
     def is_cached(self, instance):
         return hasattr(instance, self.cache_name)
 
+    def set_eri(self, qs, instances):
+        set_eri(qs,
+                self.related.field.name,
+                self.related.field.rel.field_name,
+                instances)
+
     def get_query_set(self, **db_hints):
         db = router.db_for_read(self.related.model, **db_hints)
         return self.related.model._base_manager.using(db)
@@ -239 +252 @@
     def get_prefetch_query_set(self, instances):
         vals = set(instance._get_pk_val() for instance in instances)
         params = {'%s__pk__in' % self.related.field.name: vals}
-        return (self.get_query_set(instance=instances[0]).filter(**params),
+        qs = self.get_query_set(instance=instances[0]).filter(**params)
+        self.set_eri(qs, instances)
+        return (qs,
                 attrgetter(self.related.field.attname),
                 lambda obj: obj._get_pk_val(),
                 True,
@@ -313 +328 @@
     def is_cached(self, instance):
         return hasattr(instance, self.cache_name)
 
+    def set_eri(self, qs, instances):
+        set_eri(qs,
+                self.field.related.get_accessor_name(),
+                self.field.attname,
+                instances
+                )
+
     def get_query_set(self, **db_hints):
         db = router.db_for_read(self.field.rel.to, **db_hints)
         rel_mgr = self.field.rel.to._default_manager
@@ -330 +352 @@
             params = {'%s__pk__in' % self.field.rel.field_name: vals}
         else:
             params = {'%s__in' % self.field.rel.field_name: vals}
-        return (self.get_query_set(instance=instances[0]).filter(**params),
+        qs = self.get_query_set(instance=instances[0]).filter(**params)
+        self.set_eri(qs, instances)
+        return (qs,
                 attrgetter(self.field.rel.field_name),
                 attrgetter(self.field.attname),
                 True,
@@ -462 +486 @@
                 }
                 self.model = rel_model
 
+            def set_eri(self, qs, instances):
+                set_eri(qs,
+                        rel_field.name,
+                        rel_field.rel.field_name,
+                        instances)
+
             def get_query_set(self):
                 try:
                     return self.instance._prefetched_objects_cache[rel_field.related_query_name()]
                 except (AttributeError, KeyError):
                     db = self._db or router.db_for_read(self.model, instance=self.instance)
-                    return super(RelatedManager, self).get_query_set().using(db).filter(**self.core_filters)
+                    qs = super(RelatedManager, self).get_query_set().using(db).filter(**self.core_filters)
+                    self.set_eri(qs, (self.instance,))
+                    return qs
 
             def get_prefetch_query_set(self, instances):
                 db = self._db or router.db_for_read(self.model, instance=instances[0])
                 query = {'%s__%s__in' % (rel_field.name, attname):
                              set(getattr(obj, attname) for obj in instances)}
                 qs = super(RelatedManager, self).get_query_set().using(db).filter(**query)
+                self.set_eri(qs, instances)
+
                 return (qs,
                         attrgetter(rel_field.get_attname()),
                         attrgetter(attname),

django/db/models/query.py

@@ -41 +41 @@
         self._for_write = False
         self._prefetch_related_lookups = []
         self._prefetch_done = False
+        self._existing_related_instances = {}
 
     ########################
     # PYTHON MAGIC METHODS #
@@ -294 +295 @@
                 obj, _ = get_cached_row(row, index_start, db, klass_info,
                                         offset=len(aggregate_select))
             else:
+                # Omit aggregates in object creation.
+                row_data = row[index_start:aggregate_start]
                 if skip:
-                    row_data = row[index_start:aggregate_start]
                     obj = model_cls(**dict(zip(init_list, row_data)))
                 else:
-                    # Omit aggregates in object creation.
-                    obj = model(*row[index_start:aggregate_start])
+                    obj = model(*row_data)
 
                 # Store the source database of the object
                 obj._state.db = db
@@ -315 +316 @@
                 for i, aggregate in enumerate(aggregate_select):
                     setattr(obj, aggregate, row[i+aggregate_start])
 
+            if self._existing_related_instances:
+                for attr, values in self._existing_related_instances.iteritems():
+                    descriptor = getattr(obj.__class__, attr)
+                    if len(values) > 1:
+                        try:
+                            attname = descriptor.field.attname
+                        except AttributeError:
+                            # If descriptor has no field attribute,
+                            # this is usually a reverse one-to-one
+                            # relation in which case obj PKs are used
+                            val = values[getattr(
+                                obj,
+                                descriptor.related.field.rel.field_name
+                                )]
+                        else:
+                            val = values[
+                                getattr(obj, attname)]
+                    else:
+                        val = values.values()[0]
+
+                    setattr(
+                        obj,
+                        descriptor.cache_name,
+                        val
+                        )
+
             yield obj
 
     def aggregate(self, *args, **kwargs):
@@ -860 +887 @@
         c = klass(model=self.model, query=query, using=self._db)
         c._for_write = self._for_write
         c._prefetch_related_lookups = self._prefetch_related_lookups[:]
+        c._existing_related_instances = dict(self._existing_related_instances)
         c.__dict__.update(kwargs)
         if setup and hasattr(c, '_setup_query'):
             c._setup_query()