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 subprocess
  14 import asyncio
  15 from pathlib import Path
  16
  17 import psutil
  18 import psycopg
  19 from psycopg import sql as pysql
  20
  21 from ..errors import UsageError
  22 from ..config import Configuration
  23 from ..db.connection import connect, get_pg_env, Connection, server_version_tuple, \
  24                             postgis_version_tuple, drop_tables, table_exists, execute_scalar
  25 from ..db.sql_preprocessor import SQLPreprocessor
  26 from ..db.query_pool import QueryPool
  27 from .exec_utils import run_osm2pgsql
  28 from ..version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
  29
  30 LOG = logging.getLogger()
  31
  32
  33 def _require_version(module: str, actual: Tuple[int, int], expected: Tuple[int, int]) -> None:
  34     """ Compares the version for the given module and raises an exception
  35         if the actual version is too old.
  36     """
  37     if actual < expected:
  38         LOG.fatal('Minimum supported version of %s is %d.%d. '
  39                   'Found version %d.%d.',
  40                   module, expected[0], expected[1], actual[0], actual[1])
  41         raise UsageError(f'{module} is too old.')
  42
  43
  44 def _require_loaded(extension_name: str, conn: Connection) -> None:
  45     """ Check that the given extension is loaded. """
  46     with conn.cursor() as cur:
  47         cur.execute('SELECT * FROM pg_extension WHERE extname = %s', (extension_name, ))
  48         if cur.rowcount <= 0:
  49             LOG.fatal('Required module %s is not loaded.', extension_name)
  50             raise UsageError(f'{extension_name} is not loaded.')
  51
  52
  53 def check_existing_database_plugins(dsn: str) -> None:
  54     """ Check that the database has the required plugins installed."""
  55     with connect(dsn) as conn:
  56         _require_version('PostgreSQL server',
  57                          server_version_tuple(conn),
  58                          POSTGRESQL_REQUIRED_VERSION)
  59         _require_version('PostGIS',
  60                          postgis_version_tuple(conn),
  61                          POSTGIS_REQUIRED_VERSION)
  62         _require_loaded('hstore', conn)
  63
  64
  65 def setup_database_skeleton(dsn: str, rouser: Optional[str] = None) -> None:
  66     """ Create a new database for Nominatim and populate it with the
  67         essential extensions.
  68
  69         The function fails when the database already exists or Postgresql or
  70         PostGIS versions are too old.
  71
  72         Uses `createdb` to create the database.
  73
  74         If 'rouser' is given, then the function also checks that the user
  75         with that given name exists.
  76
  77         Requires superuser rights by the caller.
  78     """
  79     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
  80
  81     if proc.returncode != 0:
  82         raise UsageError('Creating new database failed.')
  83
  84     with connect(dsn) as conn:
  85         _require_version('PostgreSQL server',
  86                          server_version_tuple(conn),
  87                          POSTGRESQL_REQUIRED_VERSION)
  88
  89         if rouser is not None:
  90             cnt = execute_scalar(conn, 'SELECT count(*) FROM pg_user where usename = %s',
  91                                  (rouser, ))
  92             if cnt == 0:
  93                 LOG.fatal("Web user '%s' does not exist. Create it with:\n"
  94                           "\n      createuser %s", rouser, rouser)
  95                 raise UsageError('Missing read-only user.')
  96
  97         # Create extensions.
  98         with conn.cursor() as cur:
  99             cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
 100             cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
 101             cur.execute('CREATE EXTENSION IF NOT EXISTS postgis_raster')
 102
 103         conn.commit()
 104
 105         _require_version('PostGIS',
 106                          postgis_version_tuple(conn),
 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 + getattr(mem, 'cached', 0)) * 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 true FROM place LIMIT 1')
 139                 if cur.rowcount == 0:
 140                     raise UsageError('No data imported by osm2pgsql.')
 141
 142         if drop:
 143             drop_tables(conn, 'planet_osm_nodes')
 144             conn.commit()
 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 table_exists(conn, '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 async def load_data(dsn: str, threads: int) -> None:
 208     """ Copy data into the word and placex table.
 209     """
 210     placex_threads = max(1, threads - 1)
 211
 212     progress = asyncio.create_task(_progress_print())
 213
 214     async with QueryPool(dsn, placex_threads + 1) as pool:
 215         # Copy data from place to placex in <threads - 1> chunks.
 216         for imod in range(placex_threads):
 217             await pool.put_query(
 218                 pysql.SQL("""INSERT INTO placex ({columns})
 219                                SELECT {columns} FROM place
 220                                 WHERE osm_id % {total} = {mod}
 221                                   AND NOT (class='place'
 222                                            and (type='houses' or type='postcode'))
 223                                   AND ST_IsValid(geometry)
 224                           """).format(columns=_COPY_COLUMNS,
 225                                       total=pysql.Literal(placex_threads),
 226                                       mod=pysql.Literal(imod)), None)
 227
 228         # Interpolations need to be copied seperately
 229         await pool.put_query("""
 230                 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' """, None)
 234
 235     progress.cancel()
 236
 237     async with await psycopg.AsyncConnection.connect(dsn) as aconn:
 238         await aconn.execute('ANALYSE')
 239
 240
 241 async def _progress_print() -> None:
 242     while True:
 243         try:
 244             await asyncio.sleep(1)
 245         except asyncio.CancelledError:
 246             print('', flush=True)
 247             break
 248         print('.', end='', flush=True)
 249
 250
 251 async def create_search_indices(conn: Connection, config: Configuration,
 252                                 drop: bool = False, threads: int = 1) -> None:
 253     """ Create tables that have explicit partitioning.
 254     """
 255
 256     # If index creation failed and left an index invalid, they need to be
 257     # cleaned out first, so that the script recreates them.
 258     with conn.cursor() as cur:
 259         cur.execute("""SELECT relname FROM pg_class, pg_index
 260                        WHERE pg_index.indisvalid = false
 261                              AND pg_index.indexrelid = pg_class.oid""")
 262         bad_indices = [row[0] for row in list(cur)]
 263         for idx in bad_indices:
 264             LOG.info("Drop invalid index %s.", idx)
 265             cur.execute(pysql.SQL('DROP INDEX {}').format(pysql.Identifier(idx)))
 266     conn.commit()
 267
 268     sql = SQLPreprocessor(conn, config)
 269
 270     await sql.run_parallel_sql_file(config.get_libpq_dsn(),
 271                                     'indices.sql', min(8, threads), drop=drop)