nominatim/api/core.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   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 Implementation of classes for API access via libraries.
   9 """
  10 from typing import Mapping, Optional, Any, AsyncIterator
  11 import asyncio
  12 import contextlib
  13 from pathlib import Path
  14
  15 import sqlalchemy as sa
  16 import sqlalchemy.ext.asyncio as sa_asyncio
  17 import asyncpg
  18
  19 from nominatim.config import Configuration
  20 from nominatim.api.status import get_status, StatusResult
  21
  22 class NominatimAPIAsync:
  23     """ API loader asynchornous version.
  24     """
  25     def __init__(self, project_dir: Path,
  26                  environ: Optional[Mapping[str, str]] = None) -> None:
  27         self.config = Configuration(project_dir, environ)
  28         self.server_version = 0
  29
  30         self._engine_lock = asyncio.Lock()
  31         self._engine: Optional[sa_asyncio.AsyncEngine] = None
  32
  33
  34     async def setup_database(self) -> None:
  35         """ Set up the engine and connection parameters.
  36
  37             This function will be implicitly called when the database is
  38             accessed for the first time. You may also call it explicitly to
  39             avoid that the first call is delayed by the setup.
  40         """
  41         async with self._engine_lock:
  42             if self._engine:
  43                 return
  44
  45             dsn = self.config.get_database_params()
  46
  47             dburl = sa.engine.URL.create(
  48                        'postgresql+asyncpg',
  49                        database=dsn.get('dbname'),
  50                        username=dsn.get('user'), password=dsn.get('password'),
  51                        host=dsn.get('host'), port=int(dsn['port']) if 'port' in dsn else None,
  52                        query={k: v for k, v in dsn.items()
  53                               if k not in ('user', 'password', 'dbname', 'host', 'port')})
  54             engine = sa_asyncio.create_async_engine(
  55                              dburl, future=True,
  56                              connect_args={'server_settings': {
  57                                 'DateStyle': 'sql,european',
  58                                 'max_parallel_workers_per_gather': '0'
  59                              }})
  60
  61             try:
  62                 async with engine.begin() as conn:
  63                     result = await conn.scalar(sa.text('SHOW server_version_num'))
  64                     self.server_version = int(result)
  65             except asyncpg.PostgresError:
  66                 self.server_version = 0
  67
  68             if self.server_version >= 110000:
  69                 @sa.event.listens_for(engine.sync_engine, "connect") # type: ignore[misc]
  70                 def _on_connect(dbapi_con: Any, _: Any) -> None:
  71                     cursor = dbapi_con.cursor()
  72                     cursor.execute("SET jit_above_cost TO '-1'")
  73                 # Make sure that all connections get the new settings
  74                 await self.close()
  75
  76             self._engine = engine
  77
  78
  79     async def close(self) -> None:
  80         """ Close all active connections to the database. The NominatimAPIAsync
  81             object remains usable after closing. If a new API functions is
  82             called, new connections are created.
  83         """
  84         if self._engine is not None:
  85             await self._engine.dispose()
  86
  87
  88     @contextlib.asynccontextmanager
  89     async def begin(self) -> AsyncIterator[sa_asyncio.AsyncConnection]:
  90         """ Create a new connection with automatic transaction handling.
  91
  92             This function may be used to get low-level access to the database.
  93             Refer to the documentation of SQLAlchemy for details how to use
  94             the connection object.
  95         """
  96         if self._engine is None:
  97             await self.setup_database()
  98
  99         assert self._engine is not None
 100
 101         async with self._engine.begin() as conn:
 102             yield conn
 103
 104
 105     async def status(self) -> StatusResult:
 106         """ Return the status of the database.
 107         """
 108         try:
 109             async with self.begin() as conn:
 110                 status = await get_status(conn)
 111         except asyncpg.PostgresError:
 112             return StatusResult(700, 'Database connection failed')
 113
 114         return status
 115
 116
 117 class NominatimAPI:
 118     """ API loader, synchronous version.
 119     """
 120
 121     def __init__(self, project_dir: Path,
 122                  environ: Optional[Mapping[str, str]] = None) -> None:
 123         self._loop = asyncio.new_event_loop()
 124         self._async_api = NominatimAPIAsync(project_dir, environ)
 125
 126
 127     def close(self) -> None:
 128         """ Close all active connections to the database. The NominatimAPIAsync
 129             object remains usable after closing. If a new API functions is
 130             called, new connections are created.
 131         """
 132         self._loop.run_until_complete(self._async_api.close())
 133         self._loop.close()
 134
 135
 136     def status(self) -> StatusResult:
 137         """ Return the status of the database.
 138         """
 139         return self._loop.run_until_complete(self._async_api.status())