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 import psycopg2.extras
  12 from psycopg2 import sql as pysql
  13
  14 from nominatim.db.connection import connect, get_pg_env
  15 from nominatim.db import utils as db_utils
  16 from nominatim.db.async_connection import DBConnection
  17 from nominatim.db.sql_preprocessor import SQLPreprocessor
  18 from nominatim.tools.exec_utils import run_osm2pgsql
  19 from nominatim.errors import UsageError
  20 from nominatim.version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
  21
  22 LOG = logging.getLogger()
  23
  24 def setup_database_skeleton(dsn, data_dir, no_partitions, rouser=None):
  25     """ Create a new database for Nominatim and populate it with the
  26         essential extensions and data.
  27     """
  28     LOG.warning('Creating database')
  29     create_db(dsn, rouser)
  30
  31     LOG.warning('Setting up database')
  32     with connect(dsn) as conn:
  33         setup_extensions(conn)
  34
  35     LOG.warning('Loading basic data')
  36     import_base_data(dsn, data_dir, no_partitions)
  37
  38
  39 def create_db(dsn, rouser=None):
  40     """ Create a new database for the given DSN. Fails when the database
  41         already exists or the PostgreSQL version is too old.
  42         Uses `createdb` to create the database.
  43
  44         If 'rouser' is given, then the function also checks that the user
  45         with that given name exists.
  46
  47         Requires superuser rights by the caller.
  48     """
  49     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
  50
  51     if proc.returncode != 0:
  52         raise UsageError('Creating new database failed.')
  53
  54     with connect(dsn) as conn:
  55         postgres_version = conn.server_version_tuple()
  56         if postgres_version < POSTGRESQL_REQUIRED_VERSION:
  57             LOG.fatal('Minimum supported version of Postgresql is %d.%d. '
  58                       'Found version %d.%d.',
  59                       POSTGRESQL_REQUIRED_VERSION[0], POSTGRESQL_REQUIRED_VERSION[1],
  60                       postgres_version[0], postgres_version[1])
  61             raise UsageError('PostgreSQL server is too old.')
  62
  63         if rouser is not None:
  64             with conn.cursor() as cur:
  65                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
  66                                  (rouser, ))
  67                 if cnt == 0:
  68                     LOG.fatal("Web user '%s' does not exists. Create it with:\n"
  69                               "\n      createuser %s", rouser, rouser)
  70                     raise UsageError('Missing read-only user.')
  71
  72
  73
  74 def setup_extensions(conn):
  75     """ Set up all extensions needed for Nominatim. Also checks that the
  76         versions of the extensions are sufficient.
  77     """
  78     with conn.cursor() as cur:
  79         cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
  80         cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
  81     conn.commit()
  82
  83     postgis_version = conn.postgis_version_tuple()
  84     if postgis_version < POSTGIS_REQUIRED_VERSION:
  85         LOG.fatal('Minimum supported version of PostGIS is %d.%d. '
  86                   'Found version %d.%d.',
  87                   POSTGIS_REQUIRED_VERSION[0], POSTGIS_REQUIRED_VERSION[1],
  88                   postgis_version[0], postgis_version[1])
  89         raise UsageError('PostGIS version is too old.')
  90
  91
  92 def import_base_data(dsn, sql_dir, ignore_partitions=False):
  93     """ Create and populate the tables with basic static data that provides
  94         the background for geocoding. Data is assumed to not yet exist.
  95     """
  96     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
  97     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
  98
  99     if ignore_partitions:
 100         with connect(dsn) as conn:
 101             with conn.cursor() as cur:
 102                 cur.execute('UPDATE country_name SET partition = 0')
 103             conn.commit()
 104
 105
 106 def import_osm_data(osm_file, options, drop=False, ignore_errors=False):
 107     """ Import the given OSM file. 'options' contains the list of
 108         default settings for osm2pgsql.
 109     """
 110     options['import_file'] = osm_file
 111     options['append'] = False
 112     options['threads'] = 1
 113
 114     if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
 115         # Make some educated guesses about cache size based on the size
 116         # of the import file and the available memory.
 117         mem = psutil.virtual_memory()
 118         fsize = os.stat(str(osm_file)).st_size
 119         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
 120                                              fsize * 2) / 1024 / 1024) + 1
 121
 122     run_osm2pgsql(options)
 123
 124     with connect(options['dsn']) as conn:
 125         if not ignore_errors:
 126             with conn.cursor() as cur:
 127                 cur.execute('SELECT * FROM place LIMIT 1')
 128                 if cur.rowcount == 0:
 129                     raise UsageError('No data imported by osm2pgsql.')
 130
 131         if drop:
 132             conn.drop_table('planet_osm_nodes')
 133
 134     if drop and options['flatnode_file']:
 135         Path(options['flatnode_file']).unlink()
 136
 137
 138 def create_tables(conn, config, reverse_only=False):
 139     """ Create the set of basic tables.
 140         When `reverse_only` is True, then the main table for searching will
 141         be skipped and only reverse search is possible.
 142     """
 143     sql = SQLPreprocessor(conn, config)
 144     sql.env.globals['db']['reverse_only'] = reverse_only
 145
 146     sql.run_sql_file(conn, 'tables.sql')
 147
 148
 149 def create_table_triggers(conn, config):
 150     """ Create the triggers for the tables. The trigger functions must already
 151         have been imported with refresh.create_functions().
 152     """
 153     sql = SQLPreprocessor(conn, config)
 154     sql.run_sql_file(conn, 'table-triggers.sql')
 155
 156
 157 def create_partition_tables(conn, config):
 158     """ Create tables that have explicit partitioning.
 159     """
 160     sql = SQLPreprocessor(conn, config)
 161     sql.run_sql_file(conn, 'partition-tables.src.sql')
 162
 163
 164 def truncate_data_tables(conn):
 165     """ Truncate all data tables to prepare for a fresh load.
 166     """
 167     with conn.cursor() as cur:
 168         cur.execute('TRUNCATE placex')
 169         cur.execute('TRUNCATE place_addressline')
 170         cur.execute('TRUNCATE location_area')
 171         cur.execute('TRUNCATE location_area_country')
 172         cur.execute('TRUNCATE location_property_tiger')
 173         cur.execute('TRUNCATE location_property_osmline')
 174         cur.execute('TRUNCATE location_postcode')
 175         if conn.table_exists('search_name'):
 176             cur.execute('TRUNCATE search_name')
 177         cur.execute('DROP SEQUENCE IF EXISTS seq_place')
 178         cur.execute('CREATE SEQUENCE seq_place start 100000')
 179
 180         cur.execute("""SELECT tablename FROM pg_tables
 181                        WHERE tablename LIKE 'location_road_%'""")
 182
 183         for table in [r[0] for r in list(cur)]:
 184             cur.execute('TRUNCATE ' + table)
 185
 186     conn.commit()
 187
 188
 189 _COPY_COLUMNS = pysql.SQL(',').join(map(pysql.Identifier,
 190                                         ('osm_type', 'osm_id', 'class', 'type',
 191                                          'name', 'admin_level', 'address',
 192                                          'extratags', 'geometry')))
 193
 194
 195 def load_data(dsn, threads):
 196     """ Copy data into the word and placex table.
 197     """
 198     sel = selectors.DefaultSelector()
 199     # Then copy data from place to placex in <threads - 1> chunks.
 200     place_threads = max(1, threads - 1)
 201     for imod in range(place_threads):
 202         conn = DBConnection(dsn)
 203         conn.connect()
 204         conn.perform(
 205             pysql.SQL("""INSERT INTO placex ({columns})
 206                            SELECT {columns} FROM place
 207                            WHERE osm_id % {total} = {mod}
 208                              AND NOT (class='place' and (type='houses' or type='postcode'))
 209                              AND ST_IsValid(geometry)
 210                       """).format(columns=_COPY_COLUMNS,
 211                                   total=pysql.Literal(place_threads),
 212                                   mod=pysql.Literal(imod)))
 213         sel.register(conn, selectors.EVENT_READ, conn)
 214
 215     # Address interpolations go into another table.
 216     conn = DBConnection(dsn)
 217     conn.connect()
 218     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
 219                       SELECT osm_id, address, geometry FROM place
 220                       WHERE class='place' and type='houses' and osm_type='W'
 221                             and ST_GeometryType(geometry) = 'ST_LineString'
 222                  """)
 223     sel.register(conn, selectors.EVENT_READ, conn)
 224
 225     # Now wait for all of them to finish.
 226     todo = place_threads + 1
 227     while todo > 0:
 228         for key, _ in sel.select(1):
 229             conn = key.data
 230             sel.unregister(conn)
 231             conn.wait()
 232             conn.close()
 233             todo -= 1
 234         print('.', end='', flush=True)
 235     print('\n')
 236
 237     with connect(dsn) as conn:
 238         with conn.cursor() as cur:
 239             cur.execute('ANALYSE')
 240
 241
 242 def create_search_indices(conn, config, drop=False):
 243     """ Create tables that have explicit partitioning.
 244     """
 245
 246     # If index creation failed and left an index invalid, they need to be
 247     # cleaned out first, so that the script recreates them.
 248     with conn.cursor() as cur:
 249         cur.execute("""SELECT relname FROM pg_class, pg_index
 250                        WHERE pg_index.indisvalid = false
 251                              AND pg_index.indexrelid = pg_class.oid""")
 252         bad_indices = [row[0] for row in list(cur)]
 253         for idx in bad_indices:
 254             LOG.info("Drop invalid index %s.", idx)
 255             cur.execute('DROP INDEX "{}"'.format(idx))
 256     conn.commit()
 257
 258     sql = SQLPreprocessor(conn, config)
 259
 260     sql.run_sql_file(conn, 'indices.sql', drop=drop)
 261
 262
 263 def create_country_names(conn, tokenizer, languages=None):
 264     """ Add default country names to search index. `languages` is a comma-
 265         separated list of language codes as used in OSM. If `languages` is not
 266         empty then only name translations for the given languages are added
 267         to the index.
 268     """
 269     if languages:
 270         languages = languages.split(',')
 271
 272     def _include_key(key):
 273         return key == 'name' or \
 274                (key.startswith('name:') and (not languages or key[5:] in languages))
 275
 276     with conn.cursor() as cur:
 277         psycopg2.extras.register_hstore(cur)
 278         cur.execute("""SELECT country_code, name FROM country_name
 279                        WHERE country_code is not null""")
 280
 281         with tokenizer.name_analyzer() as analyzer:
 282             for code, name in cur:
 283                 names = {'countrycode': code}
 284                 if code == 'gb':
 285                     names['short_name'] = 'UK'
 286                 if code == 'us':
 287                     names['short_name'] = 'United States'
 288
 289                 # country names (only in languages as provided)
 290                 if name:
 291                     names.update(((k, v) for k, v in name.items() if _include_key(k)))
 292
 293                 analyzer.add_country_names(code, names)
 294
 295     conn.commit()