forked from hans/Nominatim
reduce from 3 to 2 packages
This commit is contained in:
Sarah Hoffmann
21
src/nominatim_api/sql/async_core_library.py
Normal file
21
src/nominatim_api/sql/async_core_library.py
Normal file
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
Import the base library to use with asynchronous SQLAlchemy.
|
||||
"""
|
||||
# pylint: disable=invalid-name
|
||||
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
PGCORE_LIB = 'psycopg'
|
||||
PGCORE_ERROR: Any = psycopg.Error
|
||||
except ModuleNotFoundError:
|
||||
import asyncpg
|
||||
PGCORE_LIB = 'asyncpg'
|
||||
PGCORE_ERROR = asyncpg.PostgresError
|
||||
@@ -13,7 +13,7 @@ from typing import Any
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
|
||||
from nominatim_core.typing import SaColumn
|
||||
from ..typing import SaColumn
|
||||
|
||||
# pylint: disable=all
|
||||
|
||||
|
||||
119
src/nominatim_api/sql/sqlalchemy_schema.py
Normal file
119
src/nominatim_api/sql/sqlalchemy_schema.py
Normal file
@@ -0,0 +1,119 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
SQLAlchemy definitions for all tables used by the frontend.
|
||||
"""
|
||||
import sqlalchemy as sa
|
||||
|
||||
from .sqlalchemy_types import Geometry, KeyValueStore, IntArray
|
||||
|
||||
#pylint: disable=too-many-instance-attributes
|
||||
class SearchTables:
|
||||
""" Data class that holds the tables of the Nominatim database.
|
||||
|
||||
This schema strictly reflects the read-access view of the database.
|
||||
Any data used for updates only will not be visible.
|
||||
"""
|
||||
|
||||
def __init__(self, meta: sa.MetaData) -> None:
|
||||
self.meta = meta
|
||||
|
||||
self.import_status = sa.Table('import_status', meta,
|
||||
sa.Column('lastimportdate', sa.DateTime(True), nullable=False),
|
||||
sa.Column('sequence_id', sa.Integer),
|
||||
sa.Column('indexed', sa.Boolean))
|
||||
|
||||
self.properties = sa.Table('nominatim_properties', meta,
|
||||
sa.Column('property', sa.Text, nullable=False),
|
||||
sa.Column('value', sa.Text))
|
||||
|
||||
self.placex = sa.Table('placex', meta,
|
||||
sa.Column('place_id', sa.BigInteger, nullable=False),
|
||||
sa.Column('parent_place_id', sa.BigInteger),
|
||||
sa.Column('linked_place_id', sa.BigInteger),
|
||||
sa.Column('importance', sa.Float),
|
||||
sa.Column('indexed_date', sa.DateTime),
|
||||
sa.Column('rank_address', sa.SmallInteger),
|
||||
sa.Column('rank_search', sa.SmallInteger),
|
||||
sa.Column('indexed_status', sa.SmallInteger),
|
||||
sa.Column('osm_type', sa.String(1), nullable=False),
|
||||
sa.Column('osm_id', sa.BigInteger, nullable=False),
|
||||
sa.Column('class', sa.Text, nullable=False, key='class_'),
|
||||
sa.Column('type', sa.Text, nullable=False),
|
||||
sa.Column('admin_level', sa.SmallInteger),
|
||||
sa.Column('name', KeyValueStore),
|
||||
sa.Column('address', KeyValueStore),
|
||||
sa.Column('extratags', KeyValueStore),
|
||||
sa.Column('geometry', Geometry, nullable=False),
|
||||
sa.Column('wikipedia', sa.Text),
|
||||
sa.Column('country_code', sa.String(2)),
|
||||
sa.Column('housenumber', sa.Text),
|
||||
sa.Column('postcode', sa.Text),
|
||||
sa.Column('centroid', Geometry))
|
||||
|
||||
self.addressline = sa.Table('place_addressline', meta,
|
||||
sa.Column('place_id', sa.BigInteger),
|
||||
sa.Column('address_place_id', sa.BigInteger),
|
||||
sa.Column('distance', sa.Float),
|
||||
sa.Column('fromarea', sa.Boolean),
|
||||
sa.Column('isaddress', sa.Boolean))
|
||||
|
||||
self.postcode = sa.Table('location_postcode', meta,
|
||||
sa.Column('place_id', sa.BigInteger),
|
||||
sa.Column('parent_place_id', sa.BigInteger),
|
||||
sa.Column('rank_search', sa.SmallInteger),
|
||||
sa.Column('rank_address', sa.SmallInteger),
|
||||
sa.Column('indexed_status', sa.SmallInteger),
|
||||
sa.Column('indexed_date', sa.DateTime),
|
||||
sa.Column('country_code', sa.String(2)),
|
||||
sa.Column('postcode', sa.Text),
|
||||
sa.Column('geometry', Geometry))
|
||||
|
||||
self.osmline = sa.Table('location_property_osmline', meta,
|
||||
sa.Column('place_id', sa.BigInteger, nullable=False),
|
||||
sa.Column('osm_id', sa.BigInteger),
|
||||
sa.Column('parent_place_id', sa.BigInteger),
|
||||
sa.Column('indexed_date', sa.DateTime),
|
||||
sa.Column('startnumber', sa.Integer),
|
||||
sa.Column('endnumber', sa.Integer),
|
||||
sa.Column('step', sa.SmallInteger),
|
||||
sa.Column('indexed_status', sa.SmallInteger),
|
||||
sa.Column('linegeo', Geometry),
|
||||
sa.Column('address', KeyValueStore),
|
||||
sa.Column('postcode', sa.Text),
|
||||
sa.Column('country_code', sa.String(2)))
|
||||
|
||||
self.country_name = sa.Table('country_name', meta,
|
||||
sa.Column('country_code', sa.String(2)),
|
||||
sa.Column('name', KeyValueStore),
|
||||
sa.Column('derived_name', KeyValueStore),
|
||||
sa.Column('partition', sa.Integer))
|
||||
|
||||
self.country_grid = sa.Table('country_osm_grid', meta,
|
||||
sa.Column('country_code', sa.String(2)),
|
||||
sa.Column('area', sa.Float),
|
||||
sa.Column('geometry', Geometry))
|
||||
|
||||
# The following tables are not necessarily present.
|
||||
self.search_name = sa.Table('search_name', meta,
|
||||
sa.Column('place_id', sa.BigInteger),
|
||||
sa.Column('importance', sa.Float),
|
||||
sa.Column('search_rank', sa.SmallInteger),
|
||||
sa.Column('address_rank', sa.SmallInteger),
|
||||
sa.Column('name_vector', IntArray),
|
||||
sa.Column('nameaddress_vector', IntArray),
|
||||
sa.Column('country_code', sa.String(2)),
|
||||
sa.Column('centroid', Geometry))
|
||||
|
||||
self.tiger = sa.Table('location_property_tiger', meta,
|
||||
sa.Column('place_id', sa.BigInteger),
|
||||
sa.Column('parent_place_id', sa.BigInteger),
|
||||
sa.Column('startnumber', sa.Integer),
|
||||
sa.Column('endnumber', sa.Integer),
|
||||
sa.Column('step', sa.SmallInteger),
|
||||
sa.Column('linegeo', Geometry),
|
||||
sa.Column('postcode', sa.Text))
|
||||
17
src/nominatim_api/sql/sqlalchemy_types/__init__.py
Normal file
17
src/nominatim_api/sql/sqlalchemy_types/__init__.py
Normal file
@@ -0,0 +1,17 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
Module with custom types for SQLAlchemy
|
||||
"""
|
||||
|
||||
# See also https://github.com/PyCQA/pylint/issues/6006
|
||||
# pylint: disable=useless-import-alias
|
||||
|
||||
from .geometry import (Geometry as Geometry)
|
||||
from .int_array import (IntArray as IntArray)
|
||||
from .key_value import (KeyValueStore as KeyValueStore)
|
||||
from .json import (Json as Json)
|
||||
308
src/nominatim_api/sql/sqlalchemy_types/geometry.py
Normal file
308
src/nominatim_api/sql/sqlalchemy_types/geometry.py
Normal file
@@ -0,0 +1,308 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
Custom types for SQLAlchemy.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
from typing import Callable, Any, cast
|
||||
import sys
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
from sqlalchemy import types
|
||||
|
||||
from ...typing import SaColumn, SaBind
|
||||
|
||||
#pylint: disable=all
|
||||
|
||||
class Geometry_DistanceSpheroid(sa.sql.expression.FunctionElement[float]):
|
||||
""" Function to compute the spherical distance in meters.
|
||||
"""
|
||||
type = sa.Float()
|
||||
name = 'Geometry_DistanceSpheroid'
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
@compiles(Geometry_DistanceSpheroid) # type: ignore[no-untyped-call, misc]
|
||||
def _default_distance_spheroid(element: Geometry_DistanceSpheroid,
|
||||
compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
return "ST_DistanceSpheroid(%s,"\
|
||||
" 'SPHEROID[\"WGS 84\",6378137,298.257223563, AUTHORITY[\"EPSG\",\"7030\"]]')"\
|
||||
% compiler.process(element.clauses, **kw)
|
||||
|
||||
|
||||
@compiles(Geometry_DistanceSpheroid, 'sqlite') # type: ignore[no-untyped-call, misc]
|
||||
def _spatialite_distance_spheroid(element: Geometry_DistanceSpheroid,
|
||||
compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
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: Geometry_IsLineLike,
|
||||
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: Geometry_IsLineLike,
|
||||
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: Geometry_IsAreaLike,
|
||||
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: Geometry_IsAreaLike,
|
||||
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: Geometry_IntersectsBbox,
|
||||
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: Geometry_IntersectsBbox,
|
||||
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: Geometry_ColumnIntersectsBbox,
|
||||
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: Geometry_ColumnIntersectsBbox,
|
||||
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: Geometry_ColumnDWithin,
|
||||
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: Geometry_ColumnDWithin,
|
||||
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]
|
||||
""" Simplified type decorator for PostGIS geometry. This type
|
||||
only supports geometries in 4326 projection.
|
||||
"""
|
||||
cache_ok = True
|
||||
|
||||
def __init__(self, subtype: str = 'Geometry'):
|
||||
self.subtype = subtype
|
||||
|
||||
|
||||
def get_col_spec(self) -> str:
|
||||
return f'GEOMETRY({self.subtype}, 4326)'
|
||||
|
||||
|
||||
def bind_processor(self, dialect: 'sa.Dialect') -> Callable[[Any], str]:
|
||||
def process(value: Any) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
|
||||
return cast(str, value.to_wkt())
|
||||
return process
|
||||
|
||||
|
||||
def result_processor(self, dialect: 'sa.Dialect', coltype: object) -> Callable[[Any], str]:
|
||||
def process(value: Any) -> str:
|
||||
assert isinstance(value, str)
|
||||
return value
|
||||
return process
|
||||
|
||||
|
||||
def column_expression(self, col: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_AsEWKB(col)
|
||||
|
||||
|
||||
def bind_expression(self, bindvalue: SaBind) -> SaColumn:
|
||||
return sa.func.ST_GeomFromText(bindvalue, sa.text('4326'), type_=self)
|
||||
|
||||
|
||||
class comparator_factory(types.UserDefinedType.Comparator): # type: ignore[type-arg]
|
||||
|
||||
def intersects(self, other: SaColumn, use_index: bool = True) -> 'sa.Operators':
|
||||
if not use_index:
|
||||
return Geometry_IntersectsBbox(sa.func.coalesce(sa.null(), self.expr), 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 Geometry_IsLineLike(self)
|
||||
|
||||
|
||||
def is_area(self) -> SaColumn:
|
||||
return Geometry_IsAreaLike(self)
|
||||
|
||||
|
||||
def within_distance(self, other: SaColumn, distance: SaColumn) -> SaColumn:
|
||||
if isinstance(self.expr, sa.Column):
|
||||
return Geometry_ColumnDWithin(self.expr, other, distance)
|
||||
|
||||
return self.ST_Distance(other) < distance
|
||||
|
||||
|
||||
def ST_Distance(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_Distance(self, other, type_=sa.Float)
|
||||
|
||||
|
||||
def ST_Contains(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_Contains(self, other, type_=sa.Boolean)
|
||||
|
||||
|
||||
def ST_CoveredBy(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_CoveredBy(self, other, type_=sa.Boolean)
|
||||
|
||||
|
||||
def ST_ClosestPoint(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.coalesce(sa.func.ST_ClosestPoint(self, other, type_=Geometry),
|
||||
other)
|
||||
|
||||
|
||||
def ST_Buffer(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_Buffer(self, other, type_=Geometry)
|
||||
|
||||
|
||||
def ST_Expand(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_Expand(self, other, type_=Geometry)
|
||||
|
||||
|
||||
def ST_Collect(self) -> SaColumn:
|
||||
return sa.func.ST_Collect(self, type_=Geometry)
|
||||
|
||||
|
||||
def ST_Centroid(self) -> SaColumn:
|
||||
return sa.func.ST_Centroid(self, type_=Geometry)
|
||||
|
||||
|
||||
def ST_LineInterpolatePoint(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_LineInterpolatePoint(self, other, type_=Geometry)
|
||||
|
||||
|
||||
def ST_LineLocatePoint(self, other: SaColumn) -> SaColumn:
|
||||
return sa.func.ST_LineLocatePoint(self, other, type_=sa.Float)
|
||||
|
||||
|
||||
def distance_spheroid(self, other: SaColumn) -> SaColumn:
|
||||
return Geometry_DistanceSpheroid(self, other)
|
||||
|
||||
|
||||
@compiles(Geometry, 'sqlite') # type: ignore[no-untyped-call]
|
||||
def get_col_spec(self, *args, **kwargs): # type: ignore[no-untyped-def]
|
||||
return 'GEOMETRY'
|
||||
|
||||
|
||||
SQLITE_FUNCTION_ALIAS = (
|
||||
('ST_AsEWKB', sa.Text, 'AsEWKB'),
|
||||
('ST_GeomFromEWKT', Geometry, 'GeomFromEWKT'),
|
||||
('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(),
|
||||
"name": func,
|
||||
"identifier": func,
|
||||
"inherit_cache": True})
|
||||
|
||||
func_templ = f"{alias}(%s)"
|
||||
|
||||
def _sqlite_impl(element: Any, compiler: Any, **kw: Any) -> Any:
|
||||
return func_templ % compiler.process(element.clauses, **kw)
|
||||
|
||||
compiles(_FuncDef, 'sqlite')(_sqlite_impl) # type: ignore[no-untyped-call]
|
||||
|
||||
for alias in SQLITE_FUNCTION_ALIAS:
|
||||
_add_function_alias(*alias)
|
||||
123
src/nominatim_api/sql/sqlalchemy_types/int_array.py
Normal file
123
src/nominatim_api/sql/sqlalchemy_types/int_array.py
Normal file
@@ -0,0 +1,123 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
Custom type for an array of integers.
|
||||
"""
|
||||
from typing import Any, List, cast, Optional
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
from sqlalchemy.dialects.postgresql import ARRAY
|
||||
|
||||
from ...typing import SaDialect, SaColumn
|
||||
|
||||
# pylint: disable=all
|
||||
|
||||
class IntList(sa.types.TypeDecorator[Any]):
|
||||
""" A list of integers saved as a text of comma-separated numbers.
|
||||
"""
|
||||
impl = sa.types.Unicode
|
||||
cache_ok = True
|
||||
|
||||
def process_bind_param(self, value: Optional[Any], dialect: 'sa.Dialect') -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
assert isinstance(value, list)
|
||||
return ','.join(map(str, value))
|
||||
|
||||
def process_result_value(self, value: Optional[Any],
|
||||
dialect: SaDialect) -> Optional[List[int]]:
|
||||
return [int(v) for v in value.split(',')] if value is not None else None
|
||||
|
||||
def copy(self, **kw: Any) -> 'IntList':
|
||||
return IntList(self.impl.length)
|
||||
|
||||
|
||||
class IntArray(sa.types.TypeDecorator[Any]):
|
||||
""" Dialect-independent list of integers.
|
||||
"""
|
||||
impl = IntList
|
||||
cache_ok = True
|
||||
|
||||
def load_dialect_impl(self, dialect: SaDialect) -> sa.types.TypeEngine[Any]:
|
||||
if dialect.name == 'postgresql':
|
||||
return ARRAY(sa.Integer()) #pylint: disable=invalid-name
|
||||
|
||||
return IntList()
|
||||
|
||||
|
||||
class comparator_factory(sa.types.UserDefinedType.Comparator): # type: ignore[type-arg]
|
||||
|
||||
def __add__(self, other: SaColumn) -> 'sa.ColumnOperators':
|
||||
""" Concate the array with the given array. If one of the
|
||||
operants is null, the value of the other will be returned.
|
||||
"""
|
||||
return ArrayCat(self.expr, other)
|
||||
|
||||
|
||||
def contains(self, other: SaColumn, **kwargs: Any) -> 'sa.ColumnOperators':
|
||||
""" Return true if the array contains all the value of the argument
|
||||
array.
|
||||
"""
|
||||
return ArrayContains(self.expr, other)
|
||||
|
||||
|
||||
|
||||
class ArrayAgg(sa.sql.functions.GenericFunction[Any]):
|
||||
""" Aggregate function to collect elements in an array.
|
||||
"""
|
||||
type = IntArray()
|
||||
identifier = 'ArrayAgg'
|
||||
name = 'array_agg'
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
@compiles(ArrayAgg, 'sqlite') # type: ignore[no-untyped-call, misc]
|
||||
def sqlite_array_agg(element: ArrayAgg, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
return "group_concat(%s, ',')" % compiler.process(element.clauses, **kw)
|
||||
|
||||
|
||||
|
||||
class ArrayContains(sa.sql.expression.FunctionElement[Any]):
|
||||
""" Function to check if an array is fully contained in another.
|
||||
"""
|
||||
name = 'ArrayContains'
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
@compiles(ArrayContains) # type: ignore[no-untyped-call, misc]
|
||||
def generic_array_contains(element: ArrayContains, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
arg1, arg2 = list(element.clauses)
|
||||
return "(%s @> %s)" % (compiler.process(arg1, **kw),
|
||||
compiler.process(arg2, **kw))
|
||||
|
||||
|
||||
@compiles(ArrayContains, 'sqlite') # type: ignore[no-untyped-call, misc]
|
||||
def sqlite_array_contains(element: ArrayContains, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
return "array_contains(%s)" % compiler.process(element.clauses, **kw)
|
||||
|
||||
|
||||
|
||||
class ArrayCat(sa.sql.expression.FunctionElement[Any]):
|
||||
""" Function to check if an array is fully contained in another.
|
||||
"""
|
||||
type = IntArray()
|
||||
identifier = 'ArrayCat'
|
||||
inherit_cache = True
|
||||
|
||||
|
||||
@compiles(ArrayCat) # type: ignore[no-untyped-call, misc]
|
||||
def generic_array_cat(element: ArrayCat, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
return "array_cat(%s)" % compiler.process(element.clauses, **kw)
|
||||
|
||||
|
||||
@compiles(ArrayCat, 'sqlite') # type: ignore[no-untyped-call, misc]
|
||||
def sqlite_array_cat(element: ArrayCat, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
arg1, arg2 = list(element.clauses)
|
||||
return "(%s || ',' || %s)" % (compiler.process(arg1, **kw), compiler.process(arg2, **kw))
|
||||
|
||||
30
src/nominatim_api/sql/sqlalchemy_types/json.py
Normal file
30
src/nominatim_api/sql/sqlalchemy_types/json.py
Normal file
@@ -0,0 +1,30 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
Common json type for different dialects.
|
||||
"""
|
||||
from typing import Any
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.dialects.sqlite import JSON as sqlite_json
|
||||
|
||||
from ...typing import SaDialect
|
||||
|
||||
# pylint: disable=all
|
||||
|
||||
class Json(sa.types.TypeDecorator[Any]):
|
||||
""" Dialect-independent type for JSON.
|
||||
"""
|
||||
impl = sa.types.JSON
|
||||
cache_ok = True
|
||||
|
||||
def load_dialect_impl(self, dialect: SaDialect) -> sa.types.TypeEngine[Any]:
|
||||
if dialect.name == 'postgresql':
|
||||
return JSONB(none_as_null=True) # type: ignore[no-untyped-call]
|
||||
|
||||
return sqlite_json(none_as_null=True)
|
||||
62
src/nominatim_api/sql/sqlalchemy_types/key_value.py
Normal file
62
src/nominatim_api/sql/sqlalchemy_types/key_value.py
Normal file
@@ -0,0 +1,62 @@
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
#
|
||||
# This file is part of Nominatim. (https://nominatim.org)
|
||||
#
|
||||
# Copyright (C) 2024 by the Nominatim developer community.
|
||||
# For a full list of authors see the git log.
|
||||
"""
|
||||
A custom type that implements a simple key-value store of strings.
|
||||
"""
|
||||
from typing import Any
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.compiler import compiles
|
||||
from sqlalchemy.dialects.postgresql import HSTORE
|
||||
from sqlalchemy.dialects.sqlite import JSON as sqlite_json
|
||||
|
||||
from ...typing import SaDialect, SaColumn
|
||||
|
||||
# pylint: disable=all
|
||||
|
||||
class KeyValueStore(sa.types.TypeDecorator[Any]):
|
||||
""" Dialect-independent type of a simple key-value store of strings.
|
||||
"""
|
||||
impl = HSTORE
|
||||
cache_ok = True
|
||||
|
||||
def load_dialect_impl(self, dialect: SaDialect) -> sa.types.TypeEngine[Any]:
|
||||
if dialect.name == 'postgresql':
|
||||
return HSTORE() # type: ignore[no-untyped-call]
|
||||
|
||||
return sqlite_json(none_as_null=True)
|
||||
|
||||
|
||||
class comparator_factory(sa.types.UserDefinedType.Comparator): # type: ignore[type-arg]
|
||||
|
||||
def merge(self, other: SaColumn) -> 'sa.Operators':
|
||||
""" Merge the values from the given KeyValueStore into this
|
||||
one, overwriting values where necessary. When the argument
|
||||
is null, nothing happens.
|
||||
"""
|
||||
return KeyValueConcat(self.expr, other)
|
||||
|
||||
|
||||
class KeyValueConcat(sa.sql.expression.FunctionElement[Any]):
|
||||
""" Return the merged key-value store from the input parameters.
|
||||
"""
|
||||
type = KeyValueStore()
|
||||
name = 'JsonConcat'
|
||||
inherit_cache = True
|
||||
|
||||
@compiles(KeyValueConcat) # type: ignore[no-untyped-call, misc]
|
||||
def default_json_concat(element: KeyValueConcat, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
arg1, arg2 = list(element.clauses)
|
||||
return "(%s || coalesce(%s, ''::hstore))" % (compiler.process(arg1, **kw), compiler.process(arg2, **kw))
|
||||
|
||||
@compiles(KeyValueConcat, 'sqlite') # type: ignore[no-untyped-call, misc]
|
||||
def sqlite_json_concat(element: KeyValueConcat, compiler: 'sa.Compiled', **kw: Any) -> str:
|
||||
arg1, arg2 = list(element.clauses)
|
||||
return "json_patch(%s, coalesce(%s, '{}'))" % (compiler.process(arg1, **kw), compiler.process(arg2, **kw))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user