Ticket #5020: query.py.diff

django/trunk/django/db/models/query.py

@@ -408 +408 @@
         else:
             return self._filter_or_exclude(None, **filter_obj)
 
-    def select_related(self, true_or_false=True, depth=0):
+    def select_related(self, *args, **kwargs):
         "Returns a new QuerySet instance with '_select_related' modified."
-        return self._clone(_select_related=true_or_false, _max_related_depth=depth)
+        true_or_false = kwargs.pop('true_or_false', True)
+        depth = kwargs.pop('depth', 0)
+        return self._clone(_select_related=true_or_false, _max_related_depth=depth, _recurse_fields=fields)
 
     def order_by(self, *field_names):
         "Returns a new QuerySet instance with the ordering changed."
@@ -819 +821 @@
             raise NotImplementedError
     raise TypeError, "Got invalid lookup_type: %s" % repr(lookup_type)
 
-def get_cached_row(klass, row, index_start, max_depth=0, cur_depth=0):
+def get_cached_row(klass, row, index_start, fields=None, max_depth=0, cur_depth=0):
     """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 cur_depth > max_depth:
+    if max_depth and cur_depth >= max_depth:
         return None
 
+    fields_to_join = get_select_related_fields(klass._meta, fields)
+
     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:
-            cached_row = get_cached_row(f.rel.to, row, index_end, max_depth, cur_depth+1)
-            if cached_row:
-                rel_obj, index_end = cached_row
-                setattr(obj, f.get_cache_name(), rel_obj)
+    for f in fields_to_join.iterkeys():
+       cached_row = get_cached_row(f.rel.to, row, index_end, fields_to_join[f], max_depth, cur_depth+1)
+       if cached_row:
+               rel_obj, index_end = cached_row
+               setattr(obj, f.get_cache_name(), rel_obj)
     return obj, index_end
 
-def fill_table_cache(opts, select, tables, where, old_prefix, cache_tables_seen, max_depth=0, cur_depth=0):
+def get_select_related_fields(opts, fields=None):
+    """
+    Helper function that returns a dictionary of fields for select_related()
+    """
+    if fields is None:
+        fields_to_join = dict([(f, None) for f in opts.fields if f.rel and not f.null])
+    else:
+        fields_for_lookup = dict([(f.name, f) for f in opts.fields if f.rel])
+        fields_to_join = dict()
+        for f in fields:
+            path = f.split(LOOKUP_SEPARATOR)
+            try:
+                fn = fields_for_lookup[path[0]]
+                if fn not in fields_to_join:
+                    fields_to_join[fn] = []
+                if len(path) > 1:
+                    fields_to_join[fn].append(LOOKUP_SEPARATOR.join(path[1:]))
+            except KeyError:
+                raise FieldDoesNotExist, '%s has no field named %s' % (opts.object_name, path[0])
+    return fields_to_join
+
+def fill_table_cache(opts, select, tables, where, old_prefix, cache_tables_seen, max_depth=0, fields=None, cur_depth=0):
     """
     Helper function that recursively populates the select, tables and where (in
     place) for select_related queries.
+
+    Implicit select_related calls on NULL fields will force an INNER JOIN currently.
     """
 
     # If we've got a max_depth set and we've exceeded that depth, bail now.
-    if max_depth and cur_depth > max_depth:
+    if max_depth and cur_depth >= max_depth:
         return None
 
+    fields_to_join = get_select_related_fields(opts, fields)
+
     qn = backend.quote_name
-    for f in opts.fields:
-        if f.rel and not f.null:
-            db_table = f.rel.to._meta.db_table
-            if db_table not in cache_tables_seen:
-                tables.append(qn(db_table))
-            else: # The table was already seen, so give it a table alias.
-                new_prefix = '%s%s' % (db_table, len(cache_tables_seen))
-                tables.append('%s %s' % (qn(db_table), qn(new_prefix)))
-                db_table = new_prefix
-            cache_tables_seen.append(db_table)
-            where.append('%s.%s = %s.%s' % \
-                (qn(old_prefix), qn(f.column), qn(db_table), qn(f.rel.get_related_field().column)))
-            select.extend(['%s.%s' % (qn(db_table), qn(f2.column)) for f2 in f.rel.to._meta.fields])
-            fill_table_cache(f.rel.to._meta, select, tables, where, db_table, cache_tables_seen, max_depth, cur_depth+1)
+    for f in fields_to_join.iterkeys():
+        db_table = f.rel.to._meta.db_table
+        if db_table not in cache_tables_seen:
+            tables.append(qn(db_table))
+        else: # The table was already seen, so give it a table alias.
+            new_prefix = '%s%s' % (db_table, len(cache_tables_seen))
+            tables.append('%s %s' % (qn(db_table), qn(new_prefix)))
+            db_table = new_prefix
+        cache_tables_seen.append(db_table)
+        where.append('%s.%s = %s.%s' % \
+            (qn(old_prefix), qn(f.column), qn(db_table), qn(f.rel.get_related_field().column)))
+        select.extend(['%s.%s' % (qn(db_table), qn(f2.column)) for f2 in f.rel.to._meta.fields])
+        fill_table_cache(f.rel.to._meta, select, tables, where, db_table, cache_tables_seen, max_depth, fields_to_join[f], cur_depth+1)
 
 def parse_lookup(kwarg_items, opts):
     # Helper function that handles converting API kwargs