"""
Extended SQLAlchemy connection class that also includes access to the schema.
"""
-from typing import Any, Mapping, Sequence, Union
+from typing import cast, Any, Mapping, Sequence, Union, Dict, Optional, Set, \
+ Awaitable, Callable, TypeVar
+import asyncio
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncConnection
+from nominatim.typing import SaFromClause
from nominatim.db.sqlalchemy_schema import SearchTables
+from nominatim.db.sqlalchemy_types import Geometry
+from nominatim.api.logging import log
+
+T = TypeVar('T')
class SearchConnection:
""" An extended SQLAlchemy connection class, that also contains
"""
def __init__(self, conn: AsyncConnection,
- tables: SearchTables) -> None:
+ tables: SearchTables,
+ properties: Dict[str, Any]) -> None:
self.connection = conn
self.t = tables # pylint: disable=invalid-name
+ self._property_cache = properties
+ self._classtables: Optional[Set[str]] = None
+ self.query_timeout: Optional[int] = None
+
+
+ def set_query_timeout(self, timeout: Optional[int]) -> None:
+ """ Set the timeout after which a query over this connection
+ is cancelled.
+ """
+ self.query_timeout = timeout
async def scalar(self, sql: sa.sql.base.Executable,
) -> Any:
""" Execute a 'scalar()' query on the connection.
"""
- return await self.connection.scalar(sql, params)
+ log().sql(self.connection, sql, params)
+ return await asyncio.wait_for(self.connection.scalar(sql, params), self.query_timeout)
- async def execute(self, sql: sa.sql.base.Executable,
+ async def execute(self, sql: 'sa.Executable',
params: Union[Mapping[str, Any], Sequence[Mapping[str, Any]], None] = None
- ) -> 'sa.engine.Result[Any]':
+ ) -> 'sa.Result[Any]':
""" Execute a 'execute()' query on the connection.
"""
- return await self.connection.execute(sql, params)
+ log().sql(self.connection, sql, params)
+ return await asyncio.wait_for(self.connection.execute(sql, params), self.query_timeout)
+
+
+ async def get_property(self, name: str, cached: bool = True) -> str:
+ """ Get a property from Nominatim's property table.
+
+ Property values are normally cached so that they are only
+ retrieved from the database when they are queried for the
+ first time with this function. Set 'cached' to False to force
+ reading the property from the database.
+
+ Raises a ValueError if the property does not exist.
+ """
+ lookup_name = f'DBPROP:{name}'
+
+ if cached and lookup_name in self._property_cache:
+ return cast(str, self._property_cache[lookup_name])
+
+ sql = sa.select(self.t.properties.c.value)\
+ .where(self.t.properties.c.property == name)
+ value = await self.connection.scalar(sql)
+
+ if value is None:
+ raise ValueError(f"Property '{name}' not found in database.")
+
+ self._property_cache[lookup_name] = cast(str, value)
+
+ return cast(str, value)
+
+
+ async def get_db_property(self, name: str) -> Any:
+ """ Get a setting from the database. At the moment, only
+ 'server_version', the version of the database software, can
+ be retrieved with this function.
+
+ Raises a ValueError if the property does not exist.
+ """
+ if name != 'server_version':
+ raise ValueError(f"DB setting '{name}' not found in database.")
+
+ return self._property_cache['DB:server_version']
+
+
+ async def get_cached_value(self, group: str, name: str,
+ factory: Callable[[], Awaitable[T]]) -> T:
+ """ Access the cache for this Nominatim instance.
+ Each cache value needs to belong to a group and have a name.
+ This function is for internal API use only.
+
+ `factory` is an async callback function that produces
+ the value if it is not already cached.
+
+ Returns the cached value or the result of factory (also caching
+ the result).
+ """
+ full_name = f'{group}:{name}'
+
+ if full_name in self._property_cache:
+ return cast(T, self._property_cache[full_name])
+
+ value = await factory()
+ self._property_cache[full_name] = value
+
+ return value
+
+
+ async def get_class_table(self, cls: str, typ: str) -> Optional[SaFromClause]:
+ """ Lookup up if there is a classtype table for the given category
+ and return a SQLAlchemy table for it, if it exists.
+ """
+ if self._classtables is None:
+ res = await self.execute(sa.text("""SELECT tablename FROM pg_tables
+ WHERE tablename LIKE 'place_classtype_%'
+ """))
+ self._classtables = {r[0] for r in res}
+
+ tablename = f"place_classtype_{cls}_{typ}"
+
+ if tablename not in self._classtables:
+ return None
+
+ if tablename in self.t.meta.tables:
+ return self.t.meta.tables[tablename]
+
+ return sa.Table(tablename, self.t.meta,
+ sa.Column('place_id', sa.BigInteger),
+ sa.Column('centroid', Geometry))