1 from itertools import izip

     2 from django.db.models.query import sql

     3 from django.db.models.fields.related import ForeignKey

     4 

     5 from django.contrib.gis.db.backend import SpatialBackend

     6 from django.contrib.gis.db.models.fields import GeometryField

     7 from django.contrib.gis.db.models.sql import aggregates as gis_aggregates_module

     8 from django.contrib.gis.db.models.sql.conversion import AreaField, DistanceField, GeomField

     9 from django.contrib.gis.db.models.sql.where import GeoWhereNode

    10 from django.contrib.gis.measure import Area, Distance

    11 

    12 # Valid GIS query types.

    13 ALL_TERMS = sql.constants.QUERY_TERMS.copy()

    14 ALL_TERMS.update(SpatialBackend.gis_terms)

    15 

    16 # Pulling out other needed constants/routines to avoid attribute lookups.

    17 TABLE_NAME = sql.constants.TABLE_NAME

    18 get_proxied_model = sql.query.get_proxied_model

    19 

    20 class GeoQuery(sql.Query):

    21     """

    22     A single spatial SQL query.

    23     """

    24     # Overridding the valid query terms.

    25     query_terms = ALL_TERMS

    26     aggregates_module = gis_aggregates_module

    27 

    28     #### Methods overridden from the base Query class ####

    29     def __init__(self, model, conn):

    30         super(GeoQuery, self).__init__(model, conn, where=GeoWhereNode)

    31         # The following attributes are customized for the GeoQuerySet.

    32         # The GeoWhereNode and SpatialBackend classes contain backend-specific

    33         # routines and functions.

    34         self.custom_select = {}

    35         self.transformed_srid = None

    36         self.extra_select_fields = {}

    37 

    38     if SpatialBackend.oracle:

    39         # Have to override this so that GeoQuery, instead of OracleQuery,

    40         # is returned when unpickling.

    41         def __reduce__(self):

    42             callable, args, data = super(GeoQuery, self).__reduce__()

    43             return (unpickle_geoquery, (), data)

    44 

    45     def clone(self, *args, **kwargs):

    46         obj = super(GeoQuery, self).clone(*args, **kwargs)

    47         # Customized selection dictionary and transformed srid flag have

    48         # to also be added to obj.

    49         obj.custom_select = self.custom_select.copy()

    50         obj.transformed_srid = self.transformed_srid

    51         obj.extra_select_fields = self.extra_select_fields.copy()

    52         return obj

    53 

    54     def get_columns(self, with_aliases=False):

    55         """

    56         Return the list of columns to use in the select statement. If no

    57         columns have been specified, returns all columns relating to fields in

    58         the model.

    59 

    60         If 'with_aliases' is true, any column names that are duplicated

    61         (without the table names) are given unique aliases. This is needed in

    62         some cases to avoid ambiguitity with nested queries.

    63 

    64         This routine is overridden from Query to handle customized selection of

    65         geometry columns.

    66         """

    67         qn = self.quote_name_unless_alias

    68         qn2 = self.connection.ops.quote_name

    69         result = ['(%s) AS %s' % (self.get_extra_select_format(alias) % col[0], qn2(alias))

    70                   for alias, col in self.extra_select.iteritems()]

    71         aliases = set(self.extra_select.keys())

    72         if with_aliases:

    73             col_aliases = aliases.copy()

    74         else:

    75             col_aliases = set()

    76         if self.select:

    77             only_load = self.deferred_to_columns()

    78             # This loop customized for GeoQuery.

    79             for col, field in izip(self.select, self.select_fields):

    80                 if isinstance(col, (list, tuple)):

    81                     alias, column = col

    82                     table = self.alias_map[alias][TABLE_NAME]

    83                     if table in only_load and col not in only_load[table]:

    84                         continue

    85                     r = self.get_field_select(field, alias, column)

    86                     if with_aliases:

    87                         if col[1] in col_aliases:

    88                             c_alias = 'Col%d' % len(col_aliases)

    89                             result.append('%s AS %s' % (r, c_alias))

    90                             aliases.add(c_alias)

    91                             col_aliases.add(c_alias)

    92                         else:

    93                             result.append('%s AS %s' % (r, qn2(col[1])))

    94                             aliases.add(r)

    95                             col_aliases.add(col[1])

    96                     else:

    97                         result.append(r)

    98                         aliases.add(r)

    99                         col_aliases.add(col[1])

   100                 else:

   101                     result.append(col.as_sql(quote_func=qn))

   102 

   103                     if hasattr(col, 'alias'):

   104                         aliases.add(col.alias)

   105                         col_aliases.add(col.alias)

   106 

   107         elif self.default_cols:

   108             cols, new_aliases = self.get_default_columns(with_aliases,

   109                     col_aliases)

   110             result.extend(cols)

   111             aliases.update(new_aliases)

   112 

   113         result.extend([

   114                 '%s%s' % (

   115                     self.get_extra_select_format(alias) % aggregate.as_sql(quote_func=qn),

   116                     alias is not None and ' AS %s' % alias or ''

   117                     )

   118                 for alias, aggregate in self.aggregate_select.items()

   119         ])

   120 

   121         # This loop customized for GeoQuery.

   122         for (table, col), field in izip(self.related_select_cols, self.related_select_fields):

   123             r = self.get_field_select(field, table, col)

   124             if with_aliases and col in col_aliases:

   125                 c_alias = 'Col%d' % len(col_aliases)

   126                 result.append('%s AS %s' % (r, c_alias))

   127                 aliases.add(c_alias)

   128                 col_aliases.add(c_alias)

   129             else:

   130                 result.append(r)

   131                 aliases.add(r)

   132                 col_aliases.add(col)

   133 

   134         self._select_aliases = aliases

   135         return result

   136 

   137     def get_default_columns(self, with_aliases=False, col_aliases=None,

   138                             start_alias=None, opts=None, as_pairs=False):

   139         """

   140         Computes the default columns for selecting every field in the base

   141         model. Will sometimes be called to pull in related models (e.g. via

   142         select_related), in which case "opts" and "start_alias" will be given

   143         to provide a starting point for the traversal.

   144 

   145         Returns a list of strings, quoted appropriately for use in SQL

   146         directly, as well as a set of aliases used in the select statement (if

   147         'as_pairs' is True, returns a list of (alias, col_name) pairs instead

   148         of strings as the first component and None as the second component).

   149 

   150         This routine is overridden from Query to handle customized selection of

   151         geometry columns.

   152         """

   153         result = []

   154         if opts is None:

   155             opts = self.model._meta

   156         aliases = set()

   157         only_load = self.deferred_to_columns()

   158         # Skip all proxy to the root proxied model

   159         proxied_model = get_proxied_model(opts)

   160 

   161         if start_alias:

   162             seen = {None: start_alias}

   163         for field, model in opts.get_fields_with_model():

   164             if start_alias:

   165                 try:

   166                     alias = seen[model]

   167                 except KeyError:

   168                     if model is proxied_model:

   169                         alias = start_alias

   170                     else:

   171                         link_field = opts.get_ancestor_link(model)

   172                         alias = self.join((start_alias, model._meta.db_table,

   173                                            link_field.column, model._meta.pk.column))

   174                     seen[model] = alias

   175             else:

   176                 # If we're starting from the base model of the queryset, the

   177                 # aliases will have already been set up in pre_sql_setup(), so

   178                 # we can save time here.

   179                 alias = self.included_inherited_models[model]

   180             table = self.alias_map[alias][TABLE_NAME]

   181             if table in only_load and field.column not in only_load[table]:

   182                 continue

   183             if as_pairs:

   184                 result.append((alias, field.column))

   185                 aliases.add(alias)

   186                 continue

   187             # This part of the function is customized for GeoQuery. We

   188             # see if there was any custom selection specified in the

   189             # dictionary, and set up the selection format appropriately.

   190             field_sel = self.get_field_select(field, alias)

   191             if with_aliases and field.column in col_aliases:

   192                 c_alias = 'Col%d' % len(col_aliases)

   193                 result.append('%s AS %s' % (field_sel, c_alias))

   194                 col_aliases.add(c_alias)

   195                 aliases.add(c_alias)

   196             else:

   197                 r = field_sel

   198                 result.append(r)

   199                 aliases.add(r)

   200                 if with_aliases:

   201                     col_aliases.add(field.column)

   202         return result, aliases

   203 

   204     def resolve_columns(self, row, fields=()):

   205         """

   206         This routine is necessary so that distances and geometries returned

   207         from extra selection SQL get resolved appropriately into Python

   208         objects.

   209         """

   210         values = []

   211         aliases = self.extra_select.keys()

   212         if self.aggregates:

   213             # If we have an aggregate annotation, must extend the aliases

   214             # so their corresponding row values are included.

   215             aliases.extend([None for i in xrange(len(self.aggregates))])

   216 

   217         # Have to set a starting row number offset that is used for

   218         # determining the correct starting row index -- needed for

   219         # doing pagination with Oracle.

   220         rn_offset = 0

   221         if SpatialBackend.oracle:

   222             if self.high_mark is not None or self.low_mark: rn_offset = 1

   223         index_start = rn_offset + len(aliases)

   224 

   225         # Converting any extra selection values (e.g., geometries and

   226         # distance objects added by GeoQuerySet methods).

   227         values = [self.convert_values(v, self.extra_select_fields.get(a, None))

   228                   for v, a in izip(row[rn_offset:index_start], aliases)]

   229         if SpatialBackend.oracle or getattr(self, 'geo_values', False):

   230             # We resolve the rest of the columns if we're on Oracle or if

   231             # the `geo_values` attribute is defined.

   232             for value, field in izip(row[index_start:], fields):

   233                 values.append(self.convert_values(value, field))

   234         else:

   235             values.extend(row[index_start:])

   236         return tuple(values)

   237 

   238     def convert_values(self, value, field):

   239         """

   240         Using the same routines that Oracle does we can convert our

   241         extra selection objects into Geometry and Distance objects.

   242         TODO: Make converted objects 'lazy' for less overhead.

   243         """

   244         if SpatialBackend.oracle:

   245             # Running through Oracle's first.

   246             value = super(GeoQuery, self).convert_values(value, field or GeomField())

   247 

   248         if isinstance(field, DistanceField):

   249             # Using the field's distance attribute, can instantiate

   250             # `Distance` with the right context.

   251             value = Distance(**{field.distance_att : value})

   252         elif isinstance(field, AreaField):

   253             value = Area(**{field.area_att : value})

   254         elif isinstance(field, (GeomField, GeometryField)) and value:

   255             value = SpatialBackend.Geometry(value)

   256         return value

   257 

   258     def resolve_aggregate(self, value, aggregate):

   259         """

   260         Overridden from GeoQuery's normalize to handle the conversion of

   261         GeoAggregate objects.

   262         """

   263         if isinstance(aggregate, self.aggregates_module.GeoAggregate):

   264             if aggregate.is_extent:

   265                 return self.aggregates_module.convert_extent(value)

   266             else:

   267                 return self.aggregates_module.convert_geom(value, aggregate.source)

   268         else:

   269             return super(GeoQuery, self).resolve_aggregate(value, aggregate)

   270 

   271     #### Routines unique to GeoQuery ####

   272     def get_extra_select_format(self, alias):

   273         sel_fmt = '%s'

   274         if alias in self.custom_select:

   275             sel_fmt = sel_fmt % self.custom_select[alias]

   276         return sel_fmt

   277 

   278     def get_field_select(self, field, alias=None, column=None):

   279         """

   280         Returns the SELECT SQL string for the given field.  Figures out

   281         if any custom selection SQL is needed for the column  The `alias`

   282         keyword may be used to manually specify the database table where

   283         the column exists, if not in the model associated with this

   284         `GeoQuery`.  Similarly, `column` may be used to specify the exact

   285         column name, rather than using the `column` attribute on `field`.

   286         """

   287         sel_fmt = self.get_select_format(field)

   288         if field in self.custom_select:

   289             field_sel = sel_fmt % self.custom_select[field]

   290         else:

   291             field_sel = sel_fmt % self._field_column(field, alias, column)

   292         return field_sel

   293 

   294     def get_select_format(self, fld):

   295         """

   296         Returns the selection format string, depending on the requirements

   297         of the spatial backend.  For example, Oracle and MySQL require custom

   298         selection formats in order to retrieve geometries in OGC WKT. For all

   299         other fields a simple '%s' format string is returned.

   300         """

   301         if SpatialBackend.select and hasattr(fld, 'geom_type'):

   302             # This allows operations to be done on fields in the SELECT,

   303             # overriding their values -- used by the Oracle and MySQL

   304             # spatial backends to get database values as WKT, and by the

   305             # `transform` method.

   306             sel_fmt = SpatialBackend.select

   307 

   308             # Because WKT doesn't contain spatial reference information,

   309             # the SRID is prefixed to the returned WKT to ensure that the

   310             # transformed geometries have an SRID different than that of the

   311             # field -- this is only used by `transform` for Oracle and

   312             # SpatiaLite backends.

   313             if self.transformed_srid and ( SpatialBackend.oracle or

   314                                            SpatialBackend.spatialite ):

   315                 sel_fmt = "'SRID=%d;'||%s" % (self.transformed_srid, sel_fmt)

   316         else:

   317             sel_fmt = '%s'

   318         return sel_fmt

   319 

   320     # Private API utilities, subject to change.

   321     def _field_column(self, field, table_alias=None, column=None):

   322         """

   323         Helper function that returns the database column for the given field.

   324         The table and column are returned (quoted) in the proper format, e.g.,

   325         `"geoapp_city"."point"`.  If `table_alias` is not specified, the

   326         database table associated with the model of this `GeoQuery` will be

   327         used.  If `column` is specified, it will be used instead of the value

   328         in `field.column`.

   329         """

   330         if table_alias is None: table_alias = self.model._meta.db_table

   331         return "%s.%s" % (self.quote_name_unless_alias(table_alias),

   332                           self.connection.ops.quote_name(column or field.column))

   333 

   334     def _geo_field(self, field_name=None):

   335         """

   336         Returns the first Geometry field encountered; or specified via the

   337         `field_name` keyword.  The `field_name` may be a string specifying

   338         the geometry field on this GeoQuery's model, or a lookup string

   339         to a geometry field via a ForeignKey relation.

   340         """

   341         if field_name is None:

   342             # Incrementing until the first geographic field is found.

   343             for fld in self.model._meta.fields:

   344                 if isinstance(fld, GeometryField): return fld

   345             return False

   346         else:

   347             # Otherwise, check by the given field name -- which may be

   348             # a lookup to a _related_ geographic field.

   349             return GeoWhereNode._check_geo_field(self.model._meta, field_name)

   350 

   351 if SpatialBackend.oracle:

   352     def unpickle_geoquery():

   353         """

   354         Utility function, called by Python's unpickling machinery, that handles

   355         unpickling of GeoQuery subclasses of OracleQuery.

   356         """

   357         return GeoQuery.__new__(GeoQuery)

   358     unpickle_geoquery.__safe_for_unpickling__ = True