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, cast
  11 import sys
  12
  13 import sqlalchemy as sa
  14 from sqlalchemy import types
  15
  16 from nominatim.typing import SaColumn, SaBind
  17
  18 #pylint: disable=all
  19
  20 class Geometry(types.UserDefinedType): # type: ignore[type-arg]
  21     """ Simplified type decorator for PostGIS geometry. This type
  22         only supports geometries in 4326 projection.
  23     """
  24     cache_ok = True
  25
  26     def __init__(self, subtype: str = 'Geometry'):
  27         self.subtype = subtype
  28
  29
  30     def get_col_spec(self) -> str:
  31         return f'GEOMETRY({self.subtype}, 4326)'
  32
  33
  34     def bind_processor(self, dialect: 'sa.Dialect') -> Callable[[Any], str]:
  35         def process(value: Any) -> str:
  36             if isinstance(value, str):
  37                 return value
  38
  39             return cast(str, value.to_wkt())
  40         return process
  41
  42
  43     def result_processor(self, dialect: 'sa.Dialect', coltype: object) -> Callable[[Any], str]:
  44         def process(value: Any) -> str:
  45             assert isinstance(value, str)
  46             return value
  47         return process
  48
  49
  50     def bind_expression(self, bindvalue: SaBind) -> SaColumn:
  51         return sa.func.ST_GeomFromText(bindvalue,
  52                                        sa.bindparam('geometry_srid', value=4326, literal_execute=True),
  53                                        type_=self)
  54
  55
  56     class comparator_factory(types.UserDefinedType.Comparator): # type: ignore[type-arg]
  57
  58         def intersects(self, other: SaColumn) -> 'sa.Operators':
  59             return self.op('&&')(other)
  60
  61         def is_line_like(self) -> SaColumn:
  62             return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_LineString',
  63                                                                        'ST_MultiLineString'))
  64
  65         def is_area(self) -> SaColumn:
  66             return sa.func.ST_GeometryType(self, type_=sa.String).in_(('ST_Polygon',
  67                                                                        'ST_MultiPolygon'))
  68
  69
  70         def ST_DWithin(self, other: SaColumn, distance: SaColumn) -> SaColumn:
  71             return sa.func.ST_DWithin(self, other, distance, type_=sa.Float)
  72
  73
  74         def ST_Distance(self, other: SaColumn) -> SaColumn:
  75             return sa.func.ST_Distance(self, other, type_=sa.Float)
  76
  77
  78         def ST_Contains(self, other: SaColumn) -> SaColumn:
  79             return sa.func.ST_Contains(self, other, type_=sa.Float)
  80
  81
  82         def ST_ClosestPoint(self, other: SaColumn) -> SaColumn:
  83             return sa.func.ST_ClosestPoint(self, other, type_=Geometry)
  84
  85
  86         def ST_Buffer(self, other: SaColumn) -> SaColumn:
  87             return sa.func.ST_Buffer(self, other, type_=Geometry)
  88
  89
  90         def ST_Expand(self, other: SaColumn) -> SaColumn:
  91             return sa.func.ST_Expand(self, other, type_=Geometry)
  92
  93
  94         def ST_Collect(self) -> SaColumn:
  95             return sa.func.ST_Collect(self, type_=Geometry)
  96
  97
  98         def ST_Centroid(self) -> SaColumn:
  99             return sa.func.ST_Centroid(self, type_=Geometry)
 100
 101
 102         def ST_LineInterpolatePoint(self, other: SaColumn) -> SaColumn:
 103             return sa.func.ST_LineInterpolatePoint(self, other, type_=Geometry)
 104
 105
 106         def ST_LineLocatePoint(self, other: SaColumn) -> SaColumn:
 107             return sa.func.ST_LineLocatePoint(self, other, type_=sa.Float)