Ticket #7126: gis_select_related.diff

django/contrib/gis/db/models/query.py

@@ +1 @@
+from itertools import izip
 from django.core.exceptions import ImproperlyConfigured
 from django.db import connection
 from django.db.models.query import sql, QuerySet, Q
@@ -53 +54 @@
         obj.ewkt = self.ewkt
         return obj
 
+    def get_columns(self, with_aliases=False):
+        """
+        Return the list of columns to use in the select statement. If no
+        columns have been specified, returns all columns relating to fields in
+        the model.
+
+        If 'with_aliases' is true, any column names that are duplicated
+        (without the table names) are given unique aliases. This is needed in
+        some cases to avoid ambiguitity with nested queries.
+
+        This routine is overridden from Query to handle customized selection of 
+        geometry columns.
+        """
+        qn = self.quote_name_unless_alias
+        result = ['(%s) AS %s' % (col, alias) for alias, col in self.extra_select.iteritems()]
+        aliases = set(self.extra_select.keys())
+        if with_aliases:
+            col_aliases = aliases.copy()
+        else:
+            col_aliases = set()
+        if self.select:
+            for col, field in izip(self.select, self.select_fields):
+                if isinstance(col, (list, tuple)):
+                    # This part customized for GeoQuery.
+                    sel_fmt = self.get_select_format(field)
+                    r = '%s.%s' % (qn(col[0]), qn(col[1]))
+                    r = sel_fmt % r
+                    if with_aliases and col[1] in col_aliases:
+                        c_alias = 'Col%d' % len(col_aliases)
+                        result.append('%s AS %s' % (r, c_alias))
+                        aliases.add(c_alias)
+                        col_aliases.add(c_alias)
+                    else:
+                        result.append(r)
+                        aliases.add(r)
+                        col_aliases.add(col[1])
+                else:
+                    result.append(col.as_sql(quote_func=qn))
+                    if hasattr(col, 'alias'):
+                        aliases.add(col.alias)
+                        col_aliases.add(col.alias)
+        elif self.default_cols:
+            cols, new_aliases = self.get_default_columns(with_aliases,
+                    col_aliases)
+            result.extend(cols)
+            aliases.update(new_aliases)
+
+        for i in xrange(len(self.related_select_cols)):
+            table, col = self.related_select_cols[i]
+            field = self.related_select_fields[i]
+
+            # This part customized for GeoQuery.
+            sel_fmt = self.get_select_format(field)
+            r = '%s.%s' % (qn(table), qn(col))
+            r = sel_fmt % r
+
+            if with_aliases and col in col_aliases:
+                c_alias = 'Col%d' % len(col_aliases)
+                result.append('%s AS %s' % (r, c_alias))
+                aliases.add(c_alias)
+                col_aliases.add(c_alias)
+            else:
+                result.append(r)
+                aliases.add(r)
+                col_aliases.add(col)
+
+        self._select_aliases = aliases
+        return result
+
     def get_default_columns(self, with_aliases=False, col_aliases=None):
         """
         Computes the default columns for selecting every field in the base
@@ -79 +149 @@
                         root_pk, model._meta.pk.column))
                 seen[model] = alias
 
-            # This part of the function is customized for GeoQuerySet. We
+            # This part of the function is customized for GeoQuery. We
             # see if there was any custom selection specified in the
             # dictionary, and set up the selection format appropriately.
             sel_fmt = self.get_select_format(field)
@@ -229 +299 @@
         extent_sql = '%s(%s)' % (EXTENT, geo_col)
 
         self.query.select = [GeomSQL(extent_sql)]
+        self.query.select_fields = [None]
         try:
             esql, params = self.query.as_sql()
         except sql.datastructures.EmptyResultSet:
@@ -365 +436 @@
 
         # Only want the union SQL to be selected.
         self.query.select = [GeomSQL(union_sql)]
+        self.query.select_fields = [GeometryField]
         try:
             usql, params = self.query.as_sql()
         except sql.datastructures.EmptyResultSet: