Ticket #18177: prefetch_related_cache.patch

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()