nominatim/tools/replication.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2022 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Functions for updating a database from a replication source.
   9 """
  10 from typing import ContextManager, MutableMapping, Any, Generator, cast, Iterator
  11 from contextlib import contextmanager
  12 import datetime as dt
  13 from enum import Enum
  14 import logging
  15 import time
  16 import types
  17 import urllib.request as urlrequest
  18
  19 import requests
  20 from nominatim.db import status
  21 from nominatim.db.connection import Connection
  22 from nominatim.tools.exec_utils import run_osm2pgsql
  23 from nominatim.errors import UsageError
  24
  25 try:
  26     from osmium.replication.server import ReplicationServer
  27     from osmium import WriteHandler
  28     from osmium import version as pyo_version
  29 except ImportError as exc:
  30     logging.getLogger().critical("pyosmium not installed. Replication functions not available.\n"
  31                                  "To install pyosmium via pip: pip3 install osmium")
  32     raise UsageError("replication tools not available") from exc
  33
  34 LOG = logging.getLogger()
  35
  36 def init_replication(conn: Connection, base_url: str,
  37                      socket_timeout: int = 60) -> None:
  38     """ Set up replication for the server at the given base URL.
  39     """
  40     LOG.info("Using replication source: %s", base_url)
  41     date = status.compute_database_date(conn)
  42
  43     # margin of error to make sure we get all data
  44     date -= dt.timedelta(hours=3)
  45
  46     with _make_replication_server(base_url, socket_timeout) as repl:
  47         seq = repl.timestamp_to_sequence(date)
  48
  49     if seq is None:
  50         LOG.fatal("Cannot reach the configured replication service '%s'.\n"
  51                   "Does the URL point to a directory containing OSM update data?",
  52                   base_url)
  53         raise UsageError("Failed to reach replication service")
  54
  55     status.set_status(conn, date=date, seq=seq)
  56
  57     LOG.warning("Updates initialised at sequence %s (%s)", seq, date)
  58
  59
  60 def check_for_updates(conn: Connection, base_url: str,
  61                       socket_timeout: int = 60) -> int:
  62     """ Check if new data is available from the replication service at the
  63         given base URL.
  64     """
  65     _, seq, _ = status.get_status(conn)
  66
  67     if seq is None:
  68         LOG.error("Replication not set up. "
  69                   "Please run 'nominatim replication --init' first.")
  70         return 254
  71
  72     with _make_replication_server(base_url, socket_timeout) as repl:
  73         state = repl.get_state_info()
  74
  75     if state is None:
  76         LOG.error("Cannot get state for URL %s.", base_url)
  77         return 253
  78
  79     if state.sequence <= seq:
  80         LOG.warning("Database is up to date.")
  81         return 2
  82
  83     LOG.warning("New data available (%i => %i).", seq, state.sequence)
  84     return 0
  85
  86 class UpdateState(Enum):
  87     """ Possible states after an update has run.
  88     """
  89
  90     UP_TO_DATE = 0
  91     MORE_PENDING = 2
  92     NO_CHANGES = 3
  93
  94
  95 def update(conn: Connection, options: MutableMapping[str, Any],
  96            socket_timeout: int = 60) -> UpdateState:
  97     """ Update database from the next batch of data. Returns the state of
  98         updates according to `UpdateState`.
  99     """
 100     startdate, startseq, indexed = status.get_status(conn)
 101
 102     if startseq is None:
 103         LOG.error("Replication not set up. "
 104                   "Please run 'nominatim replication --init' first.")
 105         raise UsageError("Replication not set up.")
 106
 107     assert startdate is not None
 108
 109     if not indexed and options['indexed_only']:
 110         LOG.info("Skipping update. There is data that needs indexing.")
 111         return UpdateState.MORE_PENDING
 112
 113     last_since_update = dt.datetime.now(dt.timezone.utc) - startdate
 114     update_interval = dt.timedelta(seconds=options['update_interval'])
 115     if last_since_update < update_interval:
 116         duration = (update_interval - last_since_update).seconds
 117         LOG.warning("Sleeping for %s sec before next update.", duration)
 118         time.sleep(duration)
 119
 120     if options['import_file'].exists():
 121         options['import_file'].unlink()
 122
 123     # Read updates into file.
 124     with _make_replication_server(options['base_url'], socket_timeout) as repl:
 125         outhandler = WriteHandler(str(options['import_file']))
 126         endseq = repl.apply_diffs(outhandler, startseq + 1,
 127                                   max_size=options['max_diff_size'] * 1024)
 128         outhandler.close()
 129
 130         if endseq is None:
 131             return UpdateState.NO_CHANGES
 132
 133         # Consume updates with osm2pgsql.
 134         options['append'] = True
 135         options['disable_jit'] = conn.server_version_tuple() >= (11, 0)
 136         run_osm2pgsql(options)
 137
 138         # Write the current status to the file
 139         endstate = repl.get_state_info(endseq)
 140         status.set_status(conn, endstate.timestamp if endstate else None,
 141                           seq=endseq, indexed=False)
 142
 143     return UpdateState.UP_TO_DATE
 144
 145
 146 def _make_replication_server(url: str, timeout: int) -> ContextManager[ReplicationServer]:
 147     """ Returns a ReplicationServer in form of a context manager.
 148
 149         Creates a light wrapper around older versions of pyosmium that did
 150         not support the context manager interface.
 151     """
 152     if hasattr(ReplicationServer, '__enter__'):
 153         # Patches the open_url function for pyosmium >= 3.2
 154         # where the socket timeout is no longer respected.
 155         def patched_open_url(self: ReplicationServer, url: urlrequest.Request) -> Any:
 156             """ Download a resource from the given URL and return a byte sequence
 157                 of the content.
 158             """
 159             get_params = {
 160                 'headers': {"User-Agent" : f"Nominatim (pyosmium/{pyo_version.pyosmium_release})"},
 161                 'timeout': timeout or None,
 162                 'stream': True
 163             }
 164
 165             if self.session is not None:
 166                 return self.session.get(url.get_full_url(), **get_params)
 167
 168             @contextmanager
 169             def _get_url_with_session() -> Iterator[requests.Response]:
 170                 with requests.Session() as session:
 171                     request = session.get(url.get_full_url(), **get_params) # type: ignore
 172                     yield request
 173
 174             return _get_url_with_session()
 175
 176         repl = ReplicationServer(url)
 177         repl.open_url = types.MethodType(patched_open_url, repl)
 178
 179         return cast(ContextManager[ReplicationServer], repl)
 180
 181     @contextmanager
 182     def get_cm() -> Generator[ReplicationServer, None, None]:
 183         yield ReplicationServer(url)
 184
 185     return get_cm()