queryset-refactor: Allow specifying of specific relations to follow in

select_related(). Refs #5020.


git-svn-id: http://code.djangoproject.com/svn/django/branches/queryset-refactor@6899 bcc190cf-cafb-0310-a4f2-bffc1f526a37

django/db/models/query.py

@@ -85,13 +85,17 @@ class _QuerySet(object):
         database.
         """
         fill_cache = self.query.select_related
+        if isinstance(fill_cache, dict):
+            requested = fill_cache
+        else:
+            requested = None
         max_depth = self.query.max_depth
         index_end = len(self.model._meta.fields)
         extra_select = self.query.extra_select.keys()
         for row in self.query.results_iter():
             if fill_cache:
-                obj, index_end = get_cached_row(klass=self.model, row=row,
+                obj, index_end = get_cached_row(self.model, row, 0, max_depth,
-                        index_start=0, max_depth=max_depth)
+                        requested=requested)
             else:
                 obj = self.model(*row[:index_end])
             for i, k in enumerate(extra_select):
@@ -298,10 +302,25 @@ class _QuerySet(object):
         else:
             return self._filter_or_exclude(None, **filter_obj)
 
-    def select_related(self, true_or_false=True, depth=0):
+    def select_related(self, *fields, **kwargs):
-        """Returns a new QuerySet instance that will select related objects."""
+        """
+        Returns a new QuerySet instance that will select related objects. If
+        fields are specified, they must be ForeignKey fields and only those
+        related objects are included in the selection.
+        """
+        depth = kwargs.pop('depth', 0)
+        # TODO: Remove this? select_related(False) isn't really useful.
+        true_or_false = kwargs.pop('true_or_false', True)
+        if kwargs:
+            raise TypeError('Unexpected keyword arguments to select_related: %s'
+                    % (kwargs.keys(),))
         obj = self._clone()
-        obj.query.select_related = true_or_false
+        if fields:
+            if depth:
+                raise TypeError('Cannot pass both "depth" and fields to select_related()')
+            obj.query.add_select_related(fields)
+        else:
+            obj.query.select_related = true_or_false
         if depth:
             obj.query.max_depth = depth
         return obj
@@ -370,7 +389,7 @@ else:
 class ValuesQuerySet(QuerySet):
     def __init__(self, *args, **kwargs):
         super(ValuesQuerySet, self).__init__(*args, **kwargs)
-        # select_related isn't supported in values().
+        # select_related isn't supported in values(). (FIXME -#3358)
         self.query.select_related = False
 
         # QuerySet.clone() will also set up the _fields attribute with the
@@ -490,18 +509,26 @@ class QOperator(Q):
 
 QOr = QAnd = QOperator
 
-def get_cached_row(klass, row, index_start, max_depth=0, cur_depth=0):
+def get_cached_row(klass, row, index_start, max_depth=0, cur_depth=0,
+        requested=None):
     """Helper function that recursively returns an object with cache filled"""
 
-    # If we've got a max_depth set and we've exceeded that depth, bail now.
+    if max_depth and requested is None and cur_depth > max_depth:
-    if max_depth and cur_depth > max_depth:
+        # We've recursed deeply enough; stop now.
         return None
 
+    restricted = requested is not None
     index_end = index_start + len(klass._meta.fields)
     obj = klass(*row[index_start:index_end])
     for f in klass._meta.fields:
-        if f.rel and not f.null:
+        if f.rel and ((not restricted and not f.null) or
-            cached_row = get_cached_row(f.rel.to, row, index_end, max_depth, cur_depth+1)
+                (restricted and f.name in requested)):
+            if restricted:
+                next = requested[f.name]
+            else:
+                next = None
+            cached_row = get_cached_row(f.rel.to, row, index_end, max_depth,
+                    cur_depth+1, next)
             if cached_row:
                 rel_obj, index_end = cached_row
                 setattr(obj, f.get_cache_name(), rel_obj)

django/db/models/sql/query.py

@@ -636,15 +636,15 @@ class Query(object):
         return alias
 
     def fill_related_selections(self, opts=None, root_alias=None, cur_depth=1,
-            used=None):
+            used=None, requested=None, restricted=None):
         """
         Fill in the information needed for a select_related query. The current
-        "depth" is measured as the number of connections away from the root
+        depth is measured as the number of connections away from the root model
-        model (cur_depth == 1 means we are looking at models with direct
+        (for example, cur_depth=1 means we are looking at models with direct
         connections to the root model).
         """
-        if self.max_depth and cur_depth > self.max_depth:
+        if not restricted and self.max_depth and cur_depth > self.max_depth:
-            # We've recursed too deeply; bail out.
+            # We've recursed far enough; bail out.
             return
         if not opts:
             opts = self.model._meta
@@ -653,8 +653,18 @@ class Query(object):
         if not used:
             used = []
 
+        # Setup for the case when only particular related fields should be
+        # included in the related selection.
+        if requested is None and restricted is not False:
+            if isinstance(self.select_related, dict):
+                requested = self.select_related
+                restricted = True
+            else:
+                restricted = False
+
         for f in opts.fields:
-            if not f.rel or f.null:
+            if (not f.rel or (restricted and f.name not in requested) or
+                    (not restricted and f.null)):
                 continue
             table = f.rel.to._meta.db_table
             alias = self.join((root_alias, table, f.column,
@@ -662,8 +672,12 @@ class Query(object):
             used.append(alias)
             self.select.extend([(alias, f2.column)
                     for f2 in f.rel.to._meta.fields])
+            if restricted:
+                next = requested.get(f.name, {})
+            else:
+                next = False
             self.fill_related_selections(f.rel.to._meta, alias, cur_depth + 1,
-                    used)
+                    used, next, restricted)
 
     def add_filter(self, filter_expr, connector=AND, negate=False):
         """
@@ -1006,6 +1020,19 @@ class Query(object):
         self.select = [select]
         self.extra_select = SortedDict()
 
+    def add_select_related(self, fields):
+        """
+        Sets up the select_related data structure so that we only select
+        certain related models (as opposed to all models, when
+        self.select_related=True).
+        """
+        field_dict = {}
+        for field in fields:
+            d = field_dict
+            for part in field.split(LOOKUP_SEP):
+                d = d.setdefault(part, {})
+        self.select_related = field_dict
+
     def execute_sql(self, result_type=MULTI):
         """
         Run the query against the database and returns the result(s). The

docs/db-api.txt

@@ -744,8 +744,8 @@ related ``Person`` *and* the related ``City``::
     p = b.author         # Hits the database.
     c = p.hometown       # Hits the database.
 
-Note that ``select_related()`` does not follow foreign keys that have
+Note that, by default, ``select_related()`` does not follow foreign keys that
-``null=True``.
+have ``null=True``.
 
 Usually, using ``select_related()`` can vastly improve performance because your
 app can avoid many database calls. However, in situations with deeply nested
@@ -762,6 +762,41 @@ follow::
 
 The ``depth`` argument is new in the Django development version.
 
+**New in Django development version:** Sometimes you only need to access
+specific models that are related to your root model, not all of the related
+models. In these cases, you can pass the related field names to
+``select_related()`` and it will only follow those relations. You can even do
+this for models that are more than one relation away by separating the field
+names with double underscores, just as for filters. For example, if we have
+thise model::
+
+    class Room(models.Model):
+        # ...
+        building = models.ForeignKey(...)
+
+    class Group(models.Model):
+        # ...
+        teacher = models.ForeignKey(...)
+        room = models.ForeignKey(Room)
+        subject = models.ForeignKey(...)
+
+...and we only needed to work with the ``room`` and ``subject`` attributes, we
+could write this::
+
+    g = Group.objects.select_related('room', 'subject')
+
+This is also valid::
+
+    g = Group.objects.select_related('room__building', 'subject')
+
+...and would also pull in the ``building`` relation.
+
+You can only refer to ``ForeignKey`` relations in the list of fields passed to
+``select_related``. You *can* refer to foreign keys that have ``null=True``
+(unlike the default ``select_related()`` call). It's an error to use both a
+list of fields and the ``depth`` parameter in the same ``select_related()``
+call, since they are conflicting options.
+
 ``extra(select=None, where=None, params=None, tables=None, order_by=None)``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

tests/modeltests/select_related/models.py

@@ -129,7 +129,7 @@ __test__ = {'API_TESTS':"""
 >>> pea.genus.family.order.klass.phylum.kingdom.domain
 <Domain: Eukaryota>
 
-# Notice: one few query than above because of depth=1
+# Notice: one fewer queries than above because of depth=1
 >>> len(db.connection.queries)
 7
 
@@ -151,6 +151,38 @@ __test__ = {'API_TESTS':"""
 >>> s.id + 10 == s.a
 True
 
-# Reset DEBUG to where we found it.
+# The optional fields passed to select_related() control which related models
+# we pull in. This allows for smaller queries and can act as an alternative
+# (or, in addition to) the depth parameter.
+
+# In the next two cases, we explicitly say to select the 'genus' and
+# 'genus.family' models, leading to the same number of queries as before.
+>>> db.reset_queries()
+>>> world = Species.objects.select_related('genus__family')
+>>> [o.genus.family for o in world]
+[<Family: Drosophilidae>, <Family: Hominidae>, <Family: Fabaceae>, <Family: Amanitacae>]
+>>> len(db.connection.queries)
+1
+
+>>> db.reset_queries()
+>>> world = Species.objects.filter(genus__name='Amanita').select_related('genus__family')
+>>> [o.genus.family.order for o in world]
+[<Order: Agaricales>]
+>>> len(db.connection.queries)
+2
+
+>>> db.reset_queries()
+>>> Species.objects.all().select_related('genus__family__order').order_by('id')[0:1].get().genus.family.order.name
+u'Diptera'
+>>> len(db.connection.queries)
+1
+
+# Specifying both "depth" and fields is an error.
+>>> Species.objects.select_related('genus__family__order', depth=4)
+Traceback (most recent call last):
+...
+TypeError: Cannot pass both "depth" and fields to select_related()
+
+# Reser DEBUG to where we found it.
 >>> settings.DEBUG = False
 """}