add a timeout for DB queries

This commit is contained in:
Sarah Hoffmann 2023-08-24 14:57:33 +02:00
parent 2762c45569
commit 06a974df36
3 changed files with 26 additions and 2 deletions

View File

@ -9,6 +9,7 @@ Extended SQLAlchemy connection class that also includes access to the schema.
"""
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
@ -34,6 +35,14 @@ class SearchConnection:
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,
@ -42,7 +51,8 @@ class SearchConnection:
""" Execute a 'scalar()' query on the connection.
"""
log().sql(self.connection, sql, params)
return await self.connection.scalar(sql, params)
async with asyncio.timeout(self.query_timeout):
return await self.connection.scalar(sql, params)
async def execute(self, sql: 'sa.Executable',
@ -51,7 +61,8 @@ class SearchConnection:
""" Execute a 'execute()' query on the connection.
"""
log().sql(self.connection, sql, params)
return await self.connection.execute(sql, params)
async with asyncio.timeout(self.query_timeout):
return await self.connection.execute(sql, params)
async def get_property(self, name: str, cached: bool = True) -> str:

View File

@ -36,6 +36,8 @@ class NominatimAPIAsync:
environ: Optional[Mapping[str, str]] = None,
loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
self.config = Configuration(project_dir, environ)
self.query_timeout = self.config.get_int('QUERY_TIMEOUT') \
if self.config.QUERY_TIMEOUT else None
self.server_version = 0
if sys.version_info >= (3, 10):
@ -128,6 +130,7 @@ class NominatimAPIAsync:
"""
try:
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
status = await get_status(conn)
except (PGCORE_ERROR, sa.exc.OperationalError):
return StatusResult(700, 'Database connection failed')
@ -142,6 +145,7 @@ class NominatimAPIAsync:
"""
details = ntyp.LookupDetails.from_kwargs(params)
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
if details.keywords:
await make_query_analyzer(conn)
return await get_detailed_place(conn, place, details)
@ -154,6 +158,7 @@ class NominatimAPIAsync:
"""
details = ntyp.LookupDetails.from_kwargs(params)
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
if details.keywords:
await make_query_analyzer(conn)
return SearchResults(filter(None,
@ -173,6 +178,7 @@ class NominatimAPIAsync:
details = ntyp.ReverseDetails.from_kwargs(params)
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
if details.keywords:
await make_query_analyzer(conn)
geocoder = ReverseGeocoder(conn, details)
@ -187,6 +193,7 @@ class NominatimAPIAsync:
raise UsageError('Nothing to search for.')
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
geocoder = ForwardGeocoder(conn, ntyp.SearchDetails.from_kwargs(params))
phrases = [Phrase(PhraseType.NONE, p.strip()) for p in query.split(',')]
return await geocoder.lookup(phrases)
@ -204,6 +211,7 @@ class NominatimAPIAsync:
""" Find an address using structured search.
"""
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
details = ntyp.SearchDetails.from_kwargs(params)
phrases: List[Phrase] = []
@ -260,6 +268,7 @@ class NominatimAPIAsync:
details = ntyp.SearchDetails.from_kwargs(params)
async with self.begin() as conn:
conn.set_query_timeout(self.query_timeout)
if near_query:
phrases = [Phrase(PhraseType.NONE, p) for p in near_query.split(',')]
else:

View File

@ -214,6 +214,10 @@ NOMINATIM_SERVE_LEGACY_URLS=yes
# of connections _per worker_.
NOMINATIM_API_POOL_SIZE=10
# Timeout is seconds after which a single query to the database is cancelled.
# When empty, then timeouts are disabled.
NOMINATIM_QUERY_TIMEOUT=60
# Search elements just within countries
# If, despite not finding a point within the static grid of countries, it
# finds a geometry of a region, do not return the geometry. Return "Unable