]> git.openstreetmap.org Git - nominatim.git/blobdiff - nominatim/db/sqlalchemy_types.py
Merge remote-tracking branch 'upstream/master'
[nominatim.git] / nominatim / db / sqlalchemy_types.py
index 8e8cc9c8aa47c90d2fc6c4e6291c62c083df21ec..a36e8c462acfce3b4cc5e730b2eb5c008f1dfa14 100644 (file)
@@ -7,6 +7,7 @@
 """
 Custom types for SQLAlchemy.
 """
+from __future__ import annotations
 from typing import Callable, Any, cast
 import sys
 
@@ -37,7 +38,132 @@ def _default_distance_spheroid(element: SaColumn,
 @compiles(Geometry_DistanceSpheroid, 'sqlite') # type: ignore[no-untyped-call, misc]
 def _spatialite_distance_spheroid(element: SaColumn,
                                   compiler: 'sa.Compiled', **kw: Any) -> str:
-    return "Distance(%s, true)" % compiler.process(element.clauses, **kw)
+    return "COALESCE(Distance(%s, true), 0.0)" % compiler.process(element.clauses, **kw)
+
+
+class Geometry_IsLineLike(sa.sql.expression.FunctionElement[Any]):
+    """ Check if the geometry is a line or multiline.
+    """
+    name = 'Geometry_IsLineLike'
+    inherit_cache = True
+
+
+@compiles(Geometry_IsLineLike) # type: ignore[no-untyped-call, misc]
+def _default_is_line_like(element: SaColumn,
+                          compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "ST_GeometryType(%s) IN ('ST_LineString', 'ST_MultiLineString')" % \
+               compiler.process(element.clauses, **kw)
+
+
+@compiles(Geometry_IsLineLike, 'sqlite') # type: ignore[no-untyped-call, misc]
+def _sqlite_is_line_like(element: SaColumn,
+                         compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "ST_GeometryType(%s) IN ('LINESTRING', 'MULTILINESTRING')" % \
+               compiler.process(element.clauses, **kw)
+
+
+class Geometry_IsAreaLike(sa.sql.expression.FunctionElement[Any]):
+    """ Check if the geometry is a polygon or multipolygon.
+    """
+    name = 'Geometry_IsLineLike'
+    inherit_cache = True
+
+
+@compiles(Geometry_IsAreaLike) # type: ignore[no-untyped-call, misc]
+def _default_is_area_like(element: SaColumn,
+                          compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "ST_GeometryType(%s) IN ('ST_Polygon', 'ST_MultiPolygon')" % \
+               compiler.process(element.clauses, **kw)
+
+
+@compiles(Geometry_IsAreaLike, 'sqlite') # type: ignore[no-untyped-call, misc]
+def _sqlite_is_area_like(element: SaColumn,
+                         compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "ST_GeometryType(%s) IN ('POLYGON', 'MULTIPOLYGON')" % \
+               compiler.process(element.clauses, **kw)
+
+
+class Geometry_IntersectsBbox(sa.sql.expression.FunctionElement[Any]):
+    """ Check if the bounding boxes of the given geometries intersect.
+    """
+    name = 'Geometry_IntersectsBbox'
+    inherit_cache = True
+
+
+@compiles(Geometry_IntersectsBbox) # type: ignore[no-untyped-call, misc]
+def _default_intersects(element: SaColumn,
+                        compiler: 'sa.Compiled', **kw: Any) -> str:
+    arg1, arg2 = list(element.clauses)
+    return "%s && %s" % (compiler.process(arg1, **kw), compiler.process(arg2, **kw))
+
+
+@compiles(Geometry_IntersectsBbox, 'sqlite') # type: ignore[no-untyped-call, misc]
+def _sqlite_intersects(element: SaColumn,
+                       compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "MbrIntersects(%s) = 1" % compiler.process(element.clauses, **kw)
+
+
+class Geometry_ColumnIntersectsBbox(sa.sql.expression.FunctionElement[Any]):
+    """ Check if the bounding box of the geometry intersects with the
+        given table column, using the spatial index for the column.
+
+        The index must exist or the query may return nothing.
+    """
+    name = 'Geometry_ColumnIntersectsBbox'
+    inherit_cache = True
+
+
+@compiles(Geometry_ColumnIntersectsBbox) # type: ignore[no-untyped-call, misc]
+def default_intersects_column(element: SaColumn,
+                              compiler: 'sa.Compiled', **kw: Any) -> str:
+    arg1, arg2 = list(element.clauses)
+    return "%s && %s" % (compiler.process(arg1, **kw), compiler.process(arg2, **kw))
+
+
+@compiles(Geometry_ColumnIntersectsBbox, 'sqlite') # type: ignore[no-untyped-call, misc]
+def spatialite_intersects_column(element: SaColumn,
+                                 compiler: 'sa.Compiled', **kw: Any) -> str:
+    arg1, arg2 = list(element.clauses)
+    return "MbrIntersects(%s, %s) = 1 and "\
+           "%s.ROWID IN (SELECT ROWID FROM SpatialIndex "\
+                        "WHERE f_table_name = '%s' AND f_geometry_column = '%s' "\
+                        "AND search_frame = %s)" %(
+              compiler.process(arg1, **kw),
+              compiler.process(arg2, **kw),
+              arg1.table.name, arg1.table.name, arg1.name,
+              compiler.process(arg2, **kw))
+
+
+class Geometry_ColumnDWithin(sa.sql.expression.FunctionElement[Any]):
+    """ Check if the geometry is within the distance of the
+        given table column, using the spatial index for the column.
+
+        The index must exist or the query may return nothing.
+    """
+    name = 'Geometry_ColumnDWithin'
+    inherit_cache = True
+
+
+@compiles(Geometry_ColumnDWithin) # type: ignore[no-untyped-call, misc]
+def default_dwithin_column(element: SaColumn,
+                           compiler: 'sa.Compiled', **kw: Any) -> str:
+    return "ST_DWithin(%s)" % compiler.process(element.clauses, **kw)
+
+@compiles(Geometry_ColumnDWithin, 'sqlite') # type: ignore[no-untyped-call, misc]
+def spatialite_dwithin_column(element: SaColumn,
+                              compiler: 'sa.Compiled', **kw: Any) -> str:
+    geom1, geom2, dist = list(element.clauses)
+    return "ST_Distance(%s, %s) < %s and "\
+           "%s.ROWID IN (SELECT ROWID FROM SpatialIndex "\
+                        "WHERE f_table_name = '%s' AND f_geometry_column = '%s' "\
+                        "AND search_frame = ST_Expand(%s, %s))" %(
+              compiler.process(geom1, **kw),
+              compiler.process(geom2, **kw),
+              compiler.process(dist, **kw),
+              geom1.table.name, geom1.table.name, geom1.name,
+              compiler.process(geom2, **kw),
+              compiler.process(dist, **kw))
+
 
 
 class Geometry(types.UserDefinedType): # type: ignore[type-arg]
@@ -81,28 +207,34 @@ class Geometry(types.UserDefinedType): # type: ignore[type-arg]
     class comparator_factory(types.UserDefinedType.Comparator): # type: ignore[type-arg]
 
         def intersects(self, other: SaColumn) -> 'sa.Operators':
-            return self.op('&&')(other)
+            if isinstance(self.expr, sa.Column):
+                return Geometry_ColumnIntersectsBbox(self.expr, other)
+
+            return Geometry_IntersectsBbox(self.expr, other)
+
 
         def is_line_like(self) -> SaColumn:
-            return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_LineString',
-                                                                       'ST_MultiLineString'))
+            return Geometry_IsLineLike(self)
+
 
         def is_area(self) -> SaColumn:
-            return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_Polygon',
-                                                                       'ST_MultiPolygon'))
+            return Geometry_IsAreaLike(self)
 
 
         def ST_DWithin(self, other: SaColumn, distance: SaColumn) -> SaColumn:
-            return sa.func.ST_DWithin(self, other, distance, type_=sa.Boolean)
+            if isinstance(self.expr, sa.Column):
+                return Geometry_ColumnDWithin(self.expr, other, distance)
+
+            return sa.func.ST_DWithin(self.expr, other, distance)
 
 
         def ST_DWithin_no_index(self, other: SaColumn, distance: SaColumn) -> SaColumn:
             return sa.func.ST_DWithin(sa.func.coalesce(sa.null(), self),
-                                      other, distance, type_=sa.Boolean)
+                                      other, distance)
 
 
         def ST_Intersects_no_index(self, other: SaColumn) -> 'sa.Operators':
-            return sa.func.coalesce(sa.null(), self).op('&&')(other)
+            return Geometry_IntersectsBbox(sa.func.coalesce(sa.null(), self), other)
 
 
         def ST_Distance(self, other: SaColumn) -> SaColumn:
@@ -118,7 +250,8 @@ class Geometry(types.UserDefinedType): # type: ignore[type-arg]
 
 
         def ST_ClosestPoint(self, other: SaColumn) -> SaColumn:
-            return sa.func.ST_ClosestPoint(self, other, type_=Geometry)
+            return sa.func.coalesce(sa.func.ST_ClosestPoint(self, other, type_=Geometry),
+                                    other)
 
 
         def ST_Buffer(self, other: SaColumn) -> SaColumn:
@@ -160,11 +293,13 @@ SQLITE_FUNCTION_ALIAS = (
     ('ST_AsGeoJSON', sa.Text, 'AsGeoJSON'),
     ('ST_AsKML', sa.Text, 'AsKML'),
     ('ST_AsSVG', sa.Text, 'AsSVG'),
+    ('ST_LineLocatePoint', sa.Float, 'ST_Line_Locate_Point'),
+    ('ST_LineInterpolatePoint', sa.Float, 'ST_Line_Interpolate_Point'),
 )
 
 def _add_function_alias(func: str, ftype: type, alias: str) -> None:
     _FuncDef = type(func, (sa.sql.functions.GenericFunction, ), {
-        "type": ftype,
+        "type": ftype(),
         "name": func,
         "identifier": func,
         "inherit_cache": True})
@@ -180,4 +315,16 @@ for alias in SQLITE_FUNCTION_ALIAS:
     _add_function_alias(*alias)
 
 
+class ST_DWithin(sa.sql.functions.GenericFunction[Any]):
+    name = 'ST_DWithin'
+    inherit_cache = True
+
 
+@compiles(ST_DWithin, 'sqlite') # type: ignore[no-untyped-call, misc]
+def default_json_array_each(element: SaColumn, compiler: 'sa.Compiled', **kw: Any) -> str:
+    geom1, geom2, dist = list(element.clauses)
+    return "(MbrIntersects(%s, ST_Expand(%s, %s)) = 1 AND ST_Distance(%s, %s) <= %s)" % (
+        compiler.process(geom1, **kw), compiler.process(geom2, **kw),
+        compiler.process(dist, **kw),
+        compiler.process(geom1, **kw), compiler.process(geom2, **kw),
+        compiler.process(dist, **kw))