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