Ticket #17001: #17001-prefetch_extensions.diff

django/contrib/contenttypes/generic.py

@@ -62 +62 @@
             # This should never happen. I love comments like this, don't you?
             raise Exception("Impossible arguments to GFK.get_content_type!")
 
-    def get_prefetch_query_set(self, instances):
+    def get_prefetch_query_set(self, instances, custom_qs=None):
+        if custom_qs is not None:
+            raise ValueError("Custom queryset can't be used for this lookup")
         # For efficiency, group the instances by content type and then do one
         # query per model
         fk_dict = defaultdict(set)
@@ -320 +322 @@
                 db = self._db or router.db_for_read(self.model, instance=self.instance)
                 return super(GenericRelatedObjectManager, self).get_query_set().using(db).filter(**self.core_filters)
 
-        def get_prefetch_query_set(self, instances):
-            db = self._db or router.db_for_read(self.model, instance=instances[0])
+        def get_prefetch_query_set(self, instances, custom_qs=None):
+            if not instances:
+                return self.model._default_manager.none()
             query = {
                 '%s__pk' % self.content_type_field_name: self.content_type.id,
                 '%s__in' % self.object_id_field_name:
                     set(obj._get_pk_val() for obj in instances)
-                }
-            qs = super(GenericRelatedObjectManager, self).get_query_set().using(db).filter(**query)
+            }
+            if custom_qs is not None:
+                qs = custom_qs.filter(**query)
+            else:
+                db = self._db or router.db_for_read(self.model, instance=instances[0])
+                qs = super(GenericRelatedObjectManager, self).get_query_set()\
+                         .using(db).filter(**query)
             return (qs,
                     attrgetter(self.object_id_field_name),
                     lambda obj: obj._get_pk_val(),
                     False,
                     self.prefetch_cache_name)
 
+
+        def all(self):
+            try:
+                return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
+            except (AttributeError, KeyError):
+                return super(GenericRelatedObjectManager, self).all()
+
+
         def add(self, *objs):
             for obj in objs:
                 if not isinstance(obj, self.model):

django/db/models/__init__.py

@@ -4 +4 @@
 from django.db.models.loading import get_apps, get_app, get_models, get_model, register_models
 from django.db.models.query import Q
 from django.db.models.expressions import F
+from django.db.models.related import R
 from django.db.models.manager import Manager
 from django.db.models.base import Model
 from django.db.models.aggregates import *
 from django.db.models.fields import *
 from django.db.models.fields.subclassing import SubfieldBase
 from django.db.models.fields.files import FileField, ImageField
-from django.db.models.fields.related import ForeignKey, OneToOneField, ManyToManyField, ManyToOneRel, ManyToManyRel, OneToOneRel
-from django.db.models.deletion import CASCADE, PROTECT, SET, SET_NULL, SET_DEFAULT, DO_NOTHING, ProtectedError
+from django.db.models.fields.related import (ForeignKey, OneToOneField,
+        ManyToManyField, ManyToOneRel, ManyToManyRel, OneToOneRel)
+from django.db.models.deletion import (CASCADE, PROTECT, SET, SET_NULL,
+        SET_DEFAULT, DO_NOTHING, ProtectedError)
 from django.db.models import signals
 from django.utils.decorators import wraps
 

django/db/models/fields/related.py

@@ -236 +236 @@
         db = router.db_for_read(self.related.model, **db_hints)
         return self.related.model._base_manager.using(db)
 
-    def get_prefetch_query_set(self, instances):
+    def get_prefetch_query_set(self, instances, custom_qs=None):
+        if custom_qs is not None:
+            # TODO: This error message is too SQLish, and might be downright
+            # wrong.
+            raise ValueError(
+                "Custom querysets can't be used for one-to-one relations")
+
         vals = set(instance._get_pk_val() for instance in instances)
         params = {'%s__pk__in' % self.related.field.name: vals}
         return (self.get_query_set(),
@@ -315 +321 @@
         else:
             return QuerySet(self.field.rel.to).using(db)
 
-    def get_prefetch_query_set(self, instances):
+    def get_prefetch_query_set(self, instances, custom_qs=None):
+        if custom_qs is not None:
+            # TODO: This error message is too SQLish, and I am not even sure
+            # this desriptor is used for m2o...
+            raise ValueError(
+                "Custom querysets can't be used for many-to-one relations")
+
         vals = set(getattr(instance, self.field.attname) for instance in instances)
         other_field = self.field.rel.get_related_field()
         if other_field.rel:
@@ -460 +472 @@
                     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)
 
-            def get_prefetch_query_set(self, instances):
-                db = self._db or router.db_for_read(self.model, instance=instances[0])
+            def get_prefetch_query_set(self, instances, custom_qs=None):
+                """
+                Return a queryset that does the bulk lookup needed
+                by prefetch_related functionality.
+                """
                 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)
+                            set(getattr(obj, attname) for obj in instances)}
+                if custom_qs is not None:
+                    qs = custom_qs.filter(**query)
+                else:
+                    db = self._db or router.db_for_read(self.model, instance=instances[0])
+                    qs = super(RelatedManager, self).get_query_set().\
+                                    using(db).filter(**query)
                 return (qs,
                         attrgetter(rel_field.get_attname()),
                         attrgetter(attname),
                         False,
                         rel_field.related_query_name())
 
+            def all(self):
+                try:
+                    return self.instance._prefetched_objects_cache[rel_field.related_query_name()]
+                except (AttributeError, KeyError):
+                    return super(RelatedManager, self).all()
+
             def add(self, *objs):
                 for obj in objs:
                     if not isinstance(obj, self.model):
@@ -542 +568 @@
                 db = self._db or router.db_for_read(self.instance.__class__, instance=self.instance)
                 return super(ManyRelatedManager, self).get_query_set().using(db)._next_is_sticky().filter(**self.core_filters)
 
-        def get_prefetch_query_set(self, instances):
+        def get_prefetch_query_set(self, instances, custom_qs=None):
             instance = instances[0]
             from django.db import connections
             db = self._db or router.db_for_read(instance.__class__, instance=instance)
-            query = {'%s__pk__in' % self.query_field_name:
-                         set(obj._get_pk_val() for obj in instances)}
-            qs = super(ManyRelatedManager, self).get_query_set().using(db)._next_is_sticky().filter(**query)
+            query = {'%s__pk__in' % self.query_field_name:
+                         set(obj._get_pk_val() for obj in instances)}
+
+            if custom_qs is not None:
+                qs = custom_qs._next_is_sticky().filter(**query)
+            else:
+                qs = (super(ManyRelatedManager, self).get_query_set().using(db)
+                      ._next_is_sticky().filter(**query))
 
             # M2M: need to annotate the query in order to get the primary model
-            # that the secondary model was actually related to. We know that
-            # there will already be a join on the join table, so we can just add
-            # the select.
+            # that the secondary model was actually related to.
+
+            # We know that there will already be a join on the join table, so we
+            # can just add the select.
 
             # For non-autocreated 'through' models, can't assume we are
             # dealing with PK values.
+
+            # TODO: This is at the wrong level of abstraction. We should not
+            # be generating SQL here, but instead maybe pass this information
+            # to the connection. NoSQL camp will have problems with this, for
+            # example.
             fk = self.through._meta.get_field(self.source_field_name)
             source_col = fk.column
             join_table = self.through._meta.db_table
-            connection = connections[db]
+            if custom_qs is not None:
+                connection = connections[custom_qs.db]
+            else:
+                connection = connections[db]
+
             qn = connection.ops.quote_name
             qs = qs.extra(select={'_prefetch_related_val':
                                       '%s.%s' % (qn(join_table), qn(source_col))})
@@ -570 +611 @@
                     False,
                     self.prefetch_cache_name)
 
+        def all(self):
+            try:
+                return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
+            except (AttributeError, KeyError):
+                return super(ManyRelatedManager, self).all()
+
         # If the ManyToMany relation has an intermediary model,
         # the add and remove methods do not exist.
         if rel.through._meta.auto_created:

django/db/models/query.py

@@ -12 +12 @@
     deferred_class_factory, InvalidQuery)
 from django.db.models.deletion import Collector
 from django.db.models import sql
+from django.db.models.related import R
 from django.utils.functional import partition
 
 # Used to control how many objects are worked with at once in some cases (e.g.
@@ -1592 +1593 @@
     query.insert_values(fields, objs, raw=raw)
     return query.get_compiler(using=using).execute_sql(return_id)
 
+def prl_to_r_objs(lookups, prefix=None):
+    """
+    This little helper function will convert a list containing R objects or
+    normal lookups into all R objects list.
+    """
+    from django.db.models.sql.constants import LOOKUP_SEP
+    if prefix is None:
+        return [isinstance(lup, R) and lup or R(lup) for lup in lookups]
+    ret = []
+    for lup in lookups:
+        if isinstance(lup, R):
+            r_obj = lup._new_prefixed(prefix)
+        else:
+            r_obj = R(prefix + LOOKUP_SEP + lup)
+        ret.append(r_obj)
+    return ret
 
 def prefetch_related_objects(result_cache, related_lookups):
     """
@@ -1604 +1621 @@
 
     if len(result_cache) == 0:
         return # nothing to do
-
+    r_objs = prl_to_r_objs(related_lookups)
     model = result_cache[0].__class__
 
     # We need to be able to dynamically add to the list of prefetch_related
     # lookups that we look up (see below).  So we need some book keeping to
     # ensure we don't do duplicate work.
-    done_lookups = set() # list of lookups like foo__bar__baz
+    seen_lookups = set() # list of lookups like foo__bar__baz
     done_queries = {}    # dictionary of things like 'foo__bar': [results]
 
     auto_lookups = [] # we add to this as we go through.
     followed_descriptors = set() # recursion protection
 
-    all_lookups = itertools.chain(related_lookups, auto_lookups)
-    for lookup in all_lookups:
-        if lookup in done_lookups:
+    # For R-objects, we have two different lookups:
+    #   - lookup: This is the related object attribute name
+    #   - lookup_refpath: This is to be used when this R-object is referenced
+    #     in chained prefetches.
+    # One way to explain these would be to say lookup is how we go forward,
+    # lookup_refpath is what happened in the past.
+
+    r_objs = itertools.chain(r_objs, auto_lookups)
+    for r_obj in r_objs:
+        if r_obj.lookup_refpath in seen_lookups:
             # We've done exactly this already, skip the whole thing
             continue
-        done_lookups.add(lookup)
+        seen_lookups.add(r_obj.lookup_refpath)
 
         # Top level, the list of objects to decorate is the the result cache
         # from the primary QuerySet. It won't be for deeper levels.
         obj_list = result_cache
 
-        attrs = lookup.split(LOOKUP_SEP)
+        attrs = r_obj.lookup.split(LOOKUP_SEP)
         for level, attr in enumerate(attrs):
             # Prepare main instances
             if len(obj_list) == 0:
@@ -1655 +1679 @@
 
             # We assume that objects retrieved are homogenous (which is the premise
             # of prefetch_related), so what applies to first object applies to all.
+            # TODO: Make sure this is really true for objects coming from generic
+            # relations.
             first_obj = obj_list[0]
-            prefetcher, descriptor, attr_found, is_fetched = get_prefetcher(first_obj, attr)
+            prefetcher, descriptor, attr_found, is_fetched = \
+                get_prefetcher(first_obj, attr)
 
             if not attr_found:
                 raise AttributeError("Cannot find '%s' on %s object, '%s' is an invalid "
                                      "parameter to prefetch_related()" %
-                                     (attr, first_obj.__class__.__name__, lookup))
+                                     (attr, first_obj.__class__.__name__,
+                                      r_obj.lookup))
 
             if level == len(attrs) - 1 and prefetcher is None:
                 # Last one, this *must* resolve to something that supports
                 # prefetching, otherwise there is no point adding it and the
                 # developer asking for it has made a mistake.
-                raise ValueError("'%s' does not resolve to a item that supports "
+                raise ValueError("'%s' does not resolve to an item that supports "
                                  "prefetching - this is an invalid parameter to "
-                                 "prefetch_related()." % lookup)
+                                 "prefetch_related()." % r_obj.lookup)
 
             if prefetcher is not None and not is_fetched:
-                # Check we didn't do this already
-                current_lookup = LOOKUP_SEP.join(attrs[0:level+1])
+                current_lookup = r_obj.get_current_lookup(level)
                 if current_lookup in done_queries:
                     obj_list = done_queries[current_lookup]
                 else:
-                    obj_list, additional_prl = prefetch_one_level(obj_list, prefetcher, attr)
+                    obj_list, additional_prl = prefetch_one_level(
+                            obj_list, prefetcher, r_obj, level)
                     # We need to ensure we don't keep adding lookups from the
                     # same relationships to stop infinite recursion. So, if we
                     # are already on an automatically added lookup, don't add
                     # the new lookups from relationships we've seen already.
-                    if not (lookup in auto_lookups and
+                    if not (r_obj in auto_lookups and
                             descriptor in followed_descriptors):
-                        for f in additional_prl:
-                            new_prl = LOOKUP_SEP.join([current_lookup, f])
-                            auto_lookups.append(new_prl)
                         done_queries[current_lookup] = obj_list
+                        additional_prl = prl_to_r_objs(additional_prl,
+                                                       current_lookup)
+                        auto_lookups.extend(additional_prl)
                     followed_descriptors.add(descriptor)
+
+            elif isinstance(getattr(obj_list[0], attr), list):
+                # The current part of the lookup relates to a r_obj.to_attr
+                # defined previous fetch. This means that obj.attr is a list
+                # of related objects, and thus we must turn the obj.attr lists
+                # into a single related object list.
+                new_list = []
+                for obj in obj_list:
+                    new_list.extend(getattr(obj, attr))
+                obj_list = new_list
             else:
                 # Either a singly related object that has already been fetched
                 # (e.g. via select_related), or hopefully some other property
                 # that doesn't support prefetching but needs to be traversed.
 
                 # We replace the current list of parent objects with that list.
+                # TODO: Check what happens if attr resolves to local field?
+                # User typoing rel_attr_id instead of rel_attr? AND there are
+                # multiple parts in the path left.
                 obj_list = [getattr(obj, attr) for obj in obj_list]
 
                 # Filter out 'None' so that we can continue with nullable
@@ -1724 +1765 @@
         try:
             rel_obj = getattr(instance, attr)
             attr_found = True
+            # If we are following a r_obj lookup path which leads us through
+            # a previous fetch with to_attr, then we might end up into a list
+            # instead of related qs. This means the objects are already
+            # fetched.
+            if isinstance(rel_obj, list):
+                is_fetched = True
         except AttributeError:
             pass
     else:
@@ -1745 +1792 @@
     return prefetcher, rel_obj_descriptor, attr_found, is_fetched
 
 
-def prefetch_one_level(instances, prefetcher, attname):
+def prefetch_one_level(instances, prefetcher, r_obj, level):
     """
     Helper function for prefetch_related_objects
 
@@ -1768 +1815 @@
     # The 'values to be matched' must be hashable as they will be used
     # in a dictionary.
 
+    # to_attr is the name of the attribute we will be fetching into, to_list
+    # is False if to_attr refers to related manager. If it refers to related
+    # manager, we will be caching in rel_manager.all(), otherwise in a list.
+    to_attr, to_list = r_obj.get_to_attr(level)
+
     rel_qs, rel_obj_attr, instance_attr, single, cache_name =\
-        prefetcher.get_prefetch_query_set(instances)
+        prefetcher.get_prefetch_query_set(instances, custom_qs=r_obj.qs if to_list else None)
+
     # We have to handle the possibility that the default manager itself added
     # prefetch_related lookups to the QuerySet we just got back. We don't want to
     # trigger the prefetch_related functionality by evaluating the query.
@@ -1791 +1844 @@
         rel_obj_cache[rel_attr_val].append(rel_obj)
 
     for obj in instances:
+        # TODO: in this case we could set the reverse attribute if the relation
+        # is o2o. Both this and the TODO below are handled by select_related
+        # in the get_cached_row iterator construction. Maybe that code could
+        # be generalized and shared.
         instance_attr_val = instance_attr(obj)
         vals = rel_obj_cache.get(instance_attr_val, [])
         if single:
@@ -1800 +1857 @@
         else:
             # Multi, attribute represents a manager with an .all() method that
             # returns a QuerySet
-            qs = getattr(obj, attname).all()
-            qs._result_cache = vals
-            # We don't want the individual qs doing prefetch_related now, since we
-            # have merged this into the current work.
-            qs._prefetch_done = True
-            obj._prefetched_objects_cache[cache_name] = qs
+            # TODO: we could set the reverse relation, so that if user does
+            # access the just fetched relation in the reverse order, we would
+            # not need to do a query. We can't do this for m2m, of course.
+            if to_list:
+                setattr(obj, to_attr, vals)
+            else:
+                # Cache in the QuerySet.all().
+                qs = getattr(obj, to_attr).all()
+                qs._result_cache = vals
+                # We don't want the individual qs doing prefetch_related now,
+                # since we have merged this into the current work.
+                qs._prefetch_done = True
+                obj._prefetched_objects_cache[cache_name] = qs
     return all_related_objects, additional_prl

django/db/models/related.py

@@ -1 +1 @@
 from django.utils.encoding import smart_unicode
 from django.db.models.fields import BLANK_CHOICE_DASH
+from django.db.models.sql.constants import LOOKUP_SEP
 
 class BoundRelatedObject(object):
     def __init__(self, related_object, field_mapping, original):
@@ -36 +37 @@
                 {'%s__isnull' % self.parent_model._meta.module_name: False})
         lst = [(x._get_pk_val(), smart_unicode(x)) for x in queryset]
         return first_choice + lst
-        
+
     def get_db_prep_lookup(self, lookup_type, value, connection, prepared=False):
         # Defer to the actual field definition for db prep
         return self.field.get_db_prep_lookup(lookup_type, value,
@@ -67 +68 @@
 
     def get_cache_name(self):
         return "_%s_cache" % self.get_accessor_name()
+
+
+# Not knowing a better place for this, I just planted R here.
+# Feel free to move this to a better place or remove this comment.
+class R(object):
+    """
+    A class used for passing options to .prefetch_related. Note that instances
+    of this class should be considered immutable.
+    """
+
+    # For R-objects, we have two different internal lookup paths:
+    #   - lookup: This is the related object attribute name
+    #   - lookup_refpath: This is to be used when this R-object is referenced
+    #     in chained prefetches.
+    # Check out the source of R-objects to see what is happening there.
+    #
+    # The difference is needed, because when we chain R-objects with to_attr
+    # defined, the lookup_path (how we got here) and lookup_refpath (how to
+    # get forward from here) will be different. For example:
+    # R('foo', to_attr='foolst') -> lookup_path = foo, that is we are going
+    # to prefetch through relation foo.
+    #
+    # If there would be another qs produced by R, the lookup_refpath would
+    # need to be 'foolst__nextpart'. Otherwise we can't distinguish between
+    # two different prefetch_related lookups to 'foo' (perhaps with custom
+    # querysets).
+    #
+    # Luckily the user does not need to know anything about this.
+
+    def __init__(self, lookup, to_attr=None, qs=None):
+        if qs is not None and not to_attr:
+            raise ValueError('When custom qs is defined, to_attr '
+                             'must also be defined')
+        self.lookup = lookup
+        self.to_attr = to_attr
+        self.qs = qs._clone() if qs is not None else None
+
+    def _new_prefixed(self, prefix):
+        """
+        _new_internal is to be used when prefetches are chained internally.
+        The returned R-object is identical to self, except lookup_path
+        is prefixed with prefix.
+        """
+        new_lookup = prefix + LOOKUP_SEP + self.lookup
+        return R(new_lookup, to_attr=self.to_attr, qs=self.qs)
+
+    def __unicode__(self):
+        return ("lookup: %s, to_attr: %s, qs: %s" %
+            (self.lookup, self.to_attr or None, self.qs))
+
+    def __repr__(self):
+        return '<%s: %s>' % (self.__class__.__name__, unicode(self))
+
+    def __eq__(self, other):
+        if isinstance(other, R):
+            return self.lookup_refpath == other.lookup_refpath
+        return False
+
+    def _lookup_refpath(self):
+        if self.to_attr is None:
+            return self.lookup
+        else:
+            path, sep, last_part = self.lookup.rpartition(LOOKUP_SEP)
+            return path + sep + self.to_attr
+    lookup_refpath = property(_lookup_refpath)
+
+    def get_current_lookup(self, level):
+        """
+        Returns the first level + 1 parts of the self.lookup_refpath
+        """
+        parts = self.lookup_refpath.split(LOOKUP_SEP)
+        return LOOKUP_SEP.join(parts[0:level + 1])
+
+    def get_to_attr(self, level):
+        """
+        Returns information about into what attribute should the results be
+        fetched, and if that attribute is related object manager, or will the
+        objects be fetched into a list.
+        """
+        parts = self.lookup_refpath.split(LOOKUP_SEP)
+        if self.to_attr is None or level < len(parts) - 1:
+            return parts[level], False
+        else:
+            return self.to_attr, True