nominatim/tools/database_import.py

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