src/nominatim_db/tools/database_import.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2024 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Functions for setting up and importing a new Nominatim database.
   9 """
  10 from typing import Tuple, Optional, Union, Sequence, MutableMapping, Any
  11 import logging
  12 import os
  13 import selectors
  14 import subprocess
  15 from pathlib import Path
  16
  17 import psutil
  18 from psycopg2 import sql as pysql
  19
  20 from ..errors import UsageError
  21 from ..config import Configuration
  22 from ..db.connection import connect, get_pg_env, Connection
  23 from ..db.async_connection import DBConnection
  24 from ..db.sql_preprocessor import SQLPreprocessor
  25 from .exec_utils import run_osm2pgsql
  26 from ..version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
  27
  28 LOG = logging.getLogger()
  29
  30 def _require_version(module: str, actual: Tuple[int, int], expected: Tuple[int, int]) -> None:
  31     """ Compares the version for the given module and raises an exception
  32         if the actual version is too old.
  33     """
  34     if actual < expected:
  35         LOG.fatal('Minimum supported version of %s is %d.%d. '
  36                   'Found version %d.%d.',
  37                   module, expected[0], expected[1], actual[0], actual[1])
  38         raise UsageError(f'{module} is too old.')
  39
  40
  41 def _require_loaded(extension_name: str, conn: Connection) -> None:
  42     """ Check that the given extension is loaded. """
  43     if not conn.extension_loaded(extension_name):
  44         LOG.fatal('Required module %s is not loaded.', extension_name)
  45         raise UsageError(f'{extension_name} is not loaded.')
  46
  47
  48 def check_existing_database_plugins(dsn: str) -> None:
  49     """ Check that the database has the required plugins installed."""
  50     with connect(dsn) as conn:
  51         _require_version('PostgreSQL server',
  52                          conn.server_version_tuple(),
  53                          POSTGRESQL_REQUIRED_VERSION)
  54         _require_version('PostGIS',
  55                          conn.postgis_version_tuple(),
  56                          POSTGIS_REQUIRED_VERSION)
  57         _require_loaded('hstore', conn)
  58
  59
  60 def setup_database_skeleton(dsn: str, rouser: Optional[str] = None) -> None:
  61     """ Create a new database for Nominatim and populate it with the
  62         essential extensions.
  63
  64         The function fails when the database already exists or Postgresql or
  65         PostGIS versions are too old.
  66
  67         Uses `createdb` to create the database.
  68
  69         If 'rouser' is given, then the function also checks that the user
  70         with that given name exists.
  71
  72         Requires superuser rights by the caller.
  73     """
  74     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
  75
  76     if proc.returncode != 0:
  77         raise UsageError('Creating new database failed.')
  78
  79     with connect(dsn) as conn:
  80         _require_version('PostgreSQL server',
  81                          conn.server_version_tuple(),
  82                          POSTGRESQL_REQUIRED_VERSION)
  83
  84         if rouser is not None:
  85             with conn.cursor() as cur:
  86                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
  87                                  (rouser, ))
  88                 if cnt == 0:
  89                     LOG.fatal("Web user '%s' does not exist. Create it with:\n"
  90                               "\n      createuser %s", rouser, rouser)
  91                     raise UsageError('Missing read-only user.')
  92
  93         # Create extensions.
  94         with conn.cursor() as cur:
  95             cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
  96             cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
  97
  98             postgis_version = conn.postgis_version_tuple()
  99             if postgis_version[0] >= 3:
 100                 cur.execute('CREATE EXTENSION IF NOT EXISTS postgis_raster')
 101
 102         conn.commit()
 103
 104         _require_version('PostGIS',
 105                          conn.postgis_version_tuple(),
 106                          POSTGIS_REQUIRED_VERSION)
 107
 108
 109 def import_osm_data(osm_files: Union[Path, Sequence[Path]],
 110                     options: MutableMapping[str, Any],
 111                     drop: bool = False, ignore_errors: bool = False) -> None:
 112     """ Import the given OSM files. 'options' contains the list of
 113         default settings for osm2pgsql.
 114     """
 115     options['import_file'] = osm_files
 116     options['append'] = False
 117     options['threads'] = 1
 118
 119     if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
 120         # Make some educated guesses about cache size based on the size
 121         # of the import file and the available memory.
 122         mem = psutil.virtual_memory()
 123         fsize = 0
 124         if isinstance(osm_files, list):
 125             for fname in osm_files:
 126                 fsize += os.stat(str(fname)).st_size
 127         else:
 128             fsize = os.stat(str(osm_files)).st_size
 129         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
 130                                              fsize * 2) / 1024 / 1024) + 1
 131
 132     run_osm2pgsql(options)
 133
 134     with connect(options['dsn']) as conn:
 135         if not ignore_errors:
 136             with conn.cursor() as cur:
 137                 cur.execute('SELECT * FROM place LIMIT 1')
 138                 if cur.rowcount == 0:
 139                     raise UsageError('No data imported by osm2pgsql.')
 140
 141         if drop:
 142             conn.drop_table('planet_osm_nodes')
 143
 144     if drop and options['flatnode_file']:
 145         Path(options['flatnode_file']).unlink()
 146
 147
 148 def create_tables(conn: Connection, config: Configuration, reverse_only: bool = False) -> None:
 149     """ Create the set of basic tables.
 150         When `reverse_only` is True, then the main table for searching will
 151         be skipped and only reverse search is possible.
 152     """
 153     sql = SQLPreprocessor(conn, config)
 154     sql.env.globals['db']['reverse_only'] = reverse_only
 155
 156     sql.run_sql_file(conn, 'tables.sql')
 157
 158
 159 def create_table_triggers(conn: Connection, config: Configuration) -> None:
 160     """ Create the triggers for the tables. The trigger functions must already
 161         have been imported with refresh.create_functions().
 162     """
 163     sql = SQLPreprocessor(conn, config)
 164     sql.run_sql_file(conn, 'table-triggers.sql')
 165
 166
 167 def create_partition_tables(conn: Connection, config: Configuration) -> None:
 168     """ Create tables that have explicit partitioning.
 169     """
 170     sql = SQLPreprocessor(conn, config)
 171     sql.run_sql_file(conn, 'partition-tables.src.sql')
 172
 173
 174 def truncate_data_tables(conn: Connection) -> None:
 175     """ Truncate all data tables to prepare for a fresh load.
 176     """
 177     with conn.cursor() as cur:
 178         cur.execute('TRUNCATE placex')
 179         cur.execute('TRUNCATE place_addressline')
 180         cur.execute('TRUNCATE location_area')
 181         cur.execute('TRUNCATE location_area_country')
 182         cur.execute('TRUNCATE location_property_tiger')
 183         cur.execute('TRUNCATE location_property_osmline')
 184         cur.execute('TRUNCATE location_postcode')
 185         if conn.table_exists('search_name'):
 186             cur.execute('TRUNCATE search_name')
 187         cur.execute('DROP SEQUENCE IF EXISTS seq_place')
 188         cur.execute('CREATE SEQUENCE seq_place start 100000')
 189
 190         cur.execute("""SELECT tablename FROM pg_tables
 191                        WHERE tablename LIKE 'location_road_%'""")
 192
 193         for table in [r[0] for r in list(cur)]:
 194             cur.execute('TRUNCATE ' + table)
 195
 196     conn.commit()
 197
 198
 199 _COPY_COLUMNS = pysql.SQL(',').join(map(pysql.Identifier,
 200                                         ('osm_type', 'osm_id', 'class', 'type',
 201                                          'name', 'admin_level', 'address',
 202                                          'extratags', 'geometry')))
 203
 204
 205 def load_data(dsn: str, threads: int) -> None:
 206     """ Copy data into the word and placex table.
 207     """
 208     sel = selectors.DefaultSelector()
 209     # Then copy data from place to placex in <threads - 1> chunks.
 210     place_threads = max(1, threads - 1)
 211     for imod in range(place_threads):
 212         conn = DBConnection(dsn)
 213         conn.connect()
 214         conn.perform(
 215             pysql.SQL("""INSERT INTO placex ({columns})
 216                            SELECT {columns} FROM place
 217                            WHERE osm_id % {total} = {mod}
 218                              AND NOT (class='place' and (type='houses' or type='postcode'))
 219                              AND ST_IsValid(geometry)
 220                       """).format(columns=_COPY_COLUMNS,
 221                                   total=pysql.Literal(place_threads),
 222                                   mod=pysql.Literal(imod)))
 223         sel.register(conn, selectors.EVENT_READ, conn)
 224
 225     # Address interpolations go into another table.
 226     conn = DBConnection(dsn)
 227     conn.connect()
 228     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
 229                       SELECT osm_id, address, geometry FROM place
 230                       WHERE class='place' and type='houses' and osm_type='W'
 231                             and ST_GeometryType(geometry) = 'ST_LineString'
 232                  """)
 233     sel.register(conn, selectors.EVENT_READ, conn)
 234
 235     # Now wait for all of them to finish.
 236     todo = place_threads + 1
 237     while todo > 0:
 238         for key, _ in sel.select(1):
 239             conn = key.data
 240             sel.unregister(conn)
 241             conn.wait()
 242             conn.close()
 243             todo -= 1
 244         print('.', end='', flush=True)
 245     print('\n')
 246
 247     with connect(dsn) as syn_conn:
 248         with syn_conn.cursor() as cur:
 249             cur.execute('ANALYSE')
 250
 251
 252 def create_search_indices(conn: Connection, config: Configuration,
 253                           drop: bool = False, threads: int = 1) -> None:
 254     """ Create tables that have explicit partitioning.
 255     """
 256
 257     # If index creation failed and left an index invalid, they need to be
 258     # cleaned out first, so that the script recreates them.
 259     with conn.cursor() as cur:
 260         cur.execute("""SELECT relname FROM pg_class, pg_index
 261                        WHERE pg_index.indisvalid = false
 262                              AND pg_index.indexrelid = pg_class.oid""")
 263         bad_indices = [row[0] for row in list(cur)]
 264         for idx in bad_indices:
 265             LOG.info("Drop invalid index %s.", idx)
 266             cur.execute(pysql.SQL('DROP INDEX {}').format(pysql.Identifier(idx)))
 267     conn.commit()
 268
 269     sql = SQLPreprocessor(conn, config)
 270
 271     sql.run_parallel_sql_file(config.get_libpq_dsn(),
 272                               'indices.sql', min(8, threads), drop=drop)