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