nominatim/db/sqlalchemy_types.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2023 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Custom types for SQLAlchemy.
   9 """
  10 from typing import Callable, Any
  11
  12 import sqlalchemy as sa
  13 import sqlalchemy.types as types
  14
  15 from nominatim.typing import SaColumn
  16
  17 class Geometry(types.UserDefinedType[Any]):
  18     """ Simplified type decorator for PostGIS geometry. This type
  19         only supports geometries in 4326 projection.
  20     """
  21     cache_ok = True
  22
  23     def __init__(self, subtype: str = 'Geometry'):
  24         self.subtype = subtype
  25
  26
  27     def get_col_spec(self) -> str:
  28         return f'GEOMETRY({self.subtype}, 4326)'
  29
  30
  31     def bind_processor(self, dialect: 'sa.Dialect') -> Callable[[Any], str]:
  32         def process(value: Any) -> str:
  33             if isinstance(value, str):
  34                 return 'SRID=4326;' + value
  35
  36             return 'SRID=4326;' + value.to_wkt()
  37         return process
  38
  39
  40     def result_processor(self, dialect: 'sa.Dialect', coltype: object) -> Callable[[Any], str]:
  41         def process(value: Any) -> str:
  42             assert isinstance(value, str)
  43             return value
  44         return process
  45
  46
  47     def bind_expression(self, bindvalue: 'sa.BindParameter[Any]') -> SaColumn:
  48         return sa.func.ST_GeomFromText(bindvalue, type_=self)
  49
  50
  51     class comparator_factory(types.UserDefinedType.Comparator):
  52
  53         def intersects(self, other: SaColumn) -> SaColumn:
  54             return self.op('&&')(other)
  55
  56         def is_line_like(self) -> SaColumn:
  57             return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_LineString',
  58                                                                        'ST_MultiLineString'))
  59
  60         def is_area(self) -> SaColumn:
  61             return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_Polygon',
  62                                                                        'ST_MultiPolygon'))
  63
  64
  65         def ST_DWithin(self, other: SaColumn, distance: SaColumn) -> SaColumn:
  66             return sa.func.ST_DWithin(self, other, distance, type_=sa.Float)
  67
  68
  69         def ST_Distance(self, other: SaColumn) -> SaColumn:
  70             return sa.func.ST_Distance(self, other, type_=sa.Float)
  71
  72
  73         def ST_Contains(self, other: SaColumn) -> SaColumn:
  74             return sa.func.ST_Contains(self, other, type_=sa.Float)
  75
  76
  77         def ST_ClosestPoint(self, other: SaColumn) -> SaColumn:
  78             return sa.func.ST_ClosestPoint(self, other, type_=Geometry)
  79
  80
  81         def ST_Buffer(self, other: SaColumn) -> SaColumn:
  82             return sa.func.ST_Buffer(self, other, type_=Geometry)
  83
  84
  85         def ST_Expand(self, other: SaColumn) -> SaColumn:
  86             return sa.func.ST_Expand(self, other, type_=Geometry)
  87
  88
  89         def ST_Collect(self) -> SaColumn:
  90             return sa.func.ST_Collect(self, type_=Geometry)
  91
  92
  93         def ST_Centroid(self) -> SaColumn:
  94             return sa.func.ST_Centroid(self, type_=Geometry)
  95
  96
  97         def ST_LineInterpolatePoint(self, other: SaColumn) -> SaColumn:
  98             return sa.func.ST_LineInterpolatePoint(self, other, type_=Geometry)
  99
 100
 101         def ST_LineLocatePoint(self, other: SaColumn) -> SaColumn:
 102             return sa.func.ST_LineLocatePoint(self, other, type_=sa.Float)