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, server_version_tuple,\
  23                             postgis_version_tuple, drop_tables, table_exists, execute_scalar
  24 from ..db.async_connection import DBConnection
  25 from ..db.sql_preprocessor import SQLPreprocessor
  26 from .exec_utils import run_osm2pgsql
  27 from ..version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
  28
  29 LOG = logging.getLogger()
  30
  31 def _require_version(module: str, actual: Tuple[int, int], expected: Tuple[int, int]) -> None:
  32     """ Compares the version for the given module and raises an exception
  33         if the actual version is too old.
  34     """
  35     if actual < expected:
  36         LOG.fatal('Minimum supported version of %s is %d.%d. '
  37                   'Found version %d.%d.',
  38                   module, expected[0], expected[1], actual[0], actual[1])
  39         raise UsageError(f'{module} is too old.')
  40
  41
  42 def _require_loaded(extension_name: str, conn: Connection) -> None:
  43     """ Check that the given extension is loaded. """
  44     with conn.cursor() as cur:
  45         cur.execute('SELECT * FROM pg_extension WHERE extname = %s', (extension_name, ))
  46         if cur.rowcount <= 0:
  47             LOG.fatal('Required module %s is not loaded.', extension_name)
  48             raise UsageError(f'{extension_name} is not loaded.')
  49
  50
  51 def check_existing_database_plugins(dsn: str) -> None:
  52     """ Check that the database has the required plugins installed."""
  53     with connect(dsn) as conn:
  54         _require_version('PostgreSQL server',
  55                          server_version_tuple(conn),
  56                          POSTGRESQL_REQUIRED_VERSION)
  57         _require_version('PostGIS',
  58                          postgis_version_tuple(conn),
  59                          POSTGIS_REQUIRED_VERSION)
  60         _require_loaded('hstore', conn)
  61
  62
  63 def setup_database_skeleton(dsn: str, rouser: Optional[str] = None) -> None:
  64     """ Create a new database for Nominatim and populate it with the
  65         essential extensions.
  66
  67         The function fails when the database already exists or Postgresql or
  68         PostGIS versions are too old.
  69
  70         Uses `createdb` to create the database.
  71
  72         If 'rouser' is given, then the function also checks that the user
  73         with that given name exists.
  74
  75         Requires superuser rights by the caller.
  76     """
  77     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
  78
  79     if proc.returncode != 0:
  80         raise UsageError('Creating new database failed.')
  81
  82     with connect(dsn) as conn:
  83         _require_version('PostgreSQL server',
  84                          server_version_tuple(conn),
  85                          POSTGRESQL_REQUIRED_VERSION)
  86
  87         if rouser is not None:
  88             cnt = execute_scalar(conn, '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 = postgis_version_tuple(conn)
 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                          postgis_version_tuple(conn),
 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             drop_tables(conn, 'planet_osm_nodes')
 145             conn.commit()
 146
 147     if drop and options['flatnode_file']:
 148         Path(options['flatnode_file']).unlink()
 149
 150
 151 def create_tables(conn: Connection, config: Configuration, reverse_only: bool = False) -> None:
 152     """ Create the set of basic tables.
 153         When `reverse_only` is True, then the main table for searching will
 154         be skipped and only reverse search is possible.
 155     """
 156     sql = SQLPreprocessor(conn, config)
 157     sql.env.globals['db']['reverse_only'] = reverse_only
 158
 159     sql.run_sql_file(conn, 'tables.sql')
 160
 161
 162 def create_table_triggers(conn: Connection, config: Configuration) -> None:
 163     """ Create the triggers for the tables. The trigger functions must already
 164         have been imported with refresh.create_functions().
 165     """
 166     sql = SQLPreprocessor(conn, config)
 167     sql.run_sql_file(conn, 'table-triggers.sql')
 168
 169
 170 def create_partition_tables(conn: Connection, config: Configuration) -> None:
 171     """ Create tables that have explicit partitioning.
 172     """
 173     sql = SQLPreprocessor(conn, config)
 174     sql.run_sql_file(conn, 'partition-tables.src.sql')
 175
 176
 177 def truncate_data_tables(conn: Connection) -> None:
 178     """ Truncate all data tables to prepare for a fresh load.
 179     """
 180     with conn.cursor() as cur:
 181         cur.execute('TRUNCATE placex')
 182         cur.execute('TRUNCATE place_addressline')
 183         cur.execute('TRUNCATE location_area')
 184         cur.execute('TRUNCATE location_area_country')
 185         cur.execute('TRUNCATE location_property_tiger')
 186         cur.execute('TRUNCATE location_property_osmline')
 187         cur.execute('TRUNCATE location_postcode')
 188         if table_exists(conn, 'search_name'):
 189             cur.execute('TRUNCATE search_name')
 190         cur.execute('DROP SEQUENCE IF EXISTS seq_place')
 191         cur.execute('CREATE SEQUENCE seq_place start 100000')
 192
 193         cur.execute("""SELECT tablename FROM pg_tables
 194                        WHERE tablename LIKE 'location_road_%'""")
 195
 196         for table in [r[0] for r in list(cur)]:
 197             cur.execute('TRUNCATE ' + table)
 198
 199     conn.commit()
 200
 201
 202 _COPY_COLUMNS = pysql.SQL(',').join(map(pysql.Identifier,
 203                                         ('osm_type', 'osm_id', 'class', 'type',
 204                                          'name', 'admin_level', 'address',
 205                                          'extratags', 'geometry')))
 206
 207
 208 def load_data(dsn: str, threads: int) -> None:
 209     """ Copy data into the word and placex table.
 210     """
 211     sel = selectors.DefaultSelector()
 212     # Then copy data from place to placex in <threads - 1> chunks.
 213     place_threads = max(1, threads - 1)
 214     for imod in range(place_threads):
 215         conn = DBConnection(dsn)
 216         conn.connect()
 217         conn.perform(
 218             pysql.SQL("""INSERT INTO placex ({columns})
 219                            SELECT {columns} FROM place
 220                            WHERE osm_id % {total} = {mod}
 221                              AND NOT (class='place' and (type='houses' or type='postcode'))
 222                              AND ST_IsValid(geometry)
 223                       """).format(columns=_COPY_COLUMNS,
 224                                   total=pysql.Literal(place_threads),
 225                                   mod=pysql.Literal(imod)))
 226         sel.register(conn, selectors.EVENT_READ, conn)
 227
 228     # Address interpolations go into another table.
 229     conn = DBConnection(dsn)
 230     conn.connect()
 231     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
 232                       SELECT osm_id, address, geometry FROM place
 233                       WHERE class='place' and type='houses' and osm_type='W'
 234                             and ST_GeometryType(geometry) = 'ST_LineString'
 235                  """)
 236     sel.register(conn, selectors.EVENT_READ, conn)
 237
 238     # Now wait for all of them to finish.
 239     todo = place_threads + 1
 240     while todo > 0:
 241         for key, _ in sel.select(1):
 242             conn = key.data
 243             sel.unregister(conn)
 244             conn.wait()
 245             conn.close()
 246             todo -= 1
 247         print('.', end='', flush=True)
 248     print('\n')
 249
 250     with connect(dsn) as syn_conn:
 251         with syn_conn.cursor() as cur:
 252             cur.execute('ANALYSE')
 253
 254
 255 def create_search_indices(conn: Connection, config: Configuration,
 256                           drop: bool = False, threads: int = 1) -> None:
 257     """ Create tables that have explicit partitioning.
 258     """
 259
 260     # If index creation failed and left an index invalid, they need to be
 261     # cleaned out first, so that the script recreates them.
 262     with conn.cursor() as cur:
 263         cur.execute("""SELECT relname FROM pg_class, pg_index
 264                        WHERE pg_index.indisvalid = false
 265                              AND pg_index.indexrelid = pg_class.oid""")
 266         bad_indices = [row[0] for row in list(cur)]
 267         for idx in bad_indices:
 268             LOG.info("Drop invalid index %s.", idx)
 269             cur.execute(pysql.SQL('DROP INDEX {}').format(pysql.Identifier(idx)))
 270     conn.commit()
 271
 272     sql = SQLPreprocessor(conn, config)
 273
 274     sql.run_parallel_sql_file(config.get_libpq_dsn(),
 275                               'indices.sql', min(8, threads), drop=drop)