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_files, options, drop=False, ignore_errors=False):
 107     """ Import the given OSM files. 'options' contains the list of
 108         default settings for osm2pgsql.
 109     """
 110     options['import_file'] = osm_files
 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 = 0
 119         if isinstance(osm_files, list):
 120             for fname in osm_files:
 121                 fsize += os.stat(str(fname)).st_size
 122         else:
 123             fsize = os.stat(str(osm_files)).st_size
 124         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
 125                                              fsize * 2) / 1024 / 1024) + 1
 126
 127     run_osm2pgsql(options)
 128
 129     with connect(options['dsn']) as conn:
 130         if not ignore_errors:
 131             with conn.cursor() as cur:
 132                 cur.execute('SELECT * FROM place LIMIT 1')
 133                 if cur.rowcount == 0:
 134                     raise UsageError('No data imported by osm2pgsql.')
 135
 136         if drop:
 137             conn.drop_table('planet_osm_nodes')
 138
 139     if drop and options['flatnode_file']:
 140         Path(options['flatnode_file']).unlink()
 141
 142
 143 def create_tables(conn, config, reverse_only=False):
 144     """ Create the set of basic tables.
 145         When `reverse_only` is True, then the main table for searching will
 146         be skipped and only reverse search is possible.
 147     """
 148     sql = SQLPreprocessor(conn, config)
 149     sql.env.globals['db']['reverse_only'] = reverse_only
 150
 151     sql.run_sql_file(conn, 'tables.sql')
 152
 153
 154 def create_table_triggers(conn, config):
 155     """ Create the triggers for the tables. The trigger functions must already
 156         have been imported with refresh.create_functions().
 157     """
 158     sql = SQLPreprocessor(conn, config)
 159     sql.run_sql_file(conn, 'table-triggers.sql')
 160
 161
 162 def create_partition_tables(conn, config):
 163     """ Create tables that have explicit partitioning.
 164     """
 165     sql = SQLPreprocessor(conn, config)
 166     sql.run_sql_file(conn, 'partition-tables.src.sql')
 167
 168
 169 def truncate_data_tables(conn):
 170     """ Truncate all data tables to prepare for a fresh load.
 171     """
 172     with conn.cursor() as cur:
 173         cur.execute('TRUNCATE placex')
 174         cur.execute('TRUNCATE place_addressline')
 175         cur.execute('TRUNCATE location_area')
 176         cur.execute('TRUNCATE location_area_country')
 177         cur.execute('TRUNCATE location_property_tiger')
 178         cur.execute('TRUNCATE location_property_osmline')
 179         cur.execute('TRUNCATE location_postcode')
 180         if conn.table_exists('search_name'):
 181             cur.execute('TRUNCATE search_name')
 182         cur.execute('DROP SEQUENCE IF EXISTS seq_place')
 183         cur.execute('CREATE SEQUENCE seq_place start 100000')
 184
 185         cur.execute("""SELECT tablename FROM pg_tables
 186                        WHERE tablename LIKE 'location_road_%'""")
 187
 188         for table in [r[0] for r in list(cur)]:
 189             cur.execute('TRUNCATE ' + table)
 190
 191     conn.commit()
 192
 193
 194 _COPY_COLUMNS = pysql.SQL(',').join(map(pysql.Identifier,
 195                                         ('osm_type', 'osm_id', 'class', 'type',
 196                                          'name', 'admin_level', 'address',
 197                                          'extratags', 'geometry')))
 198
 199
 200 def load_data(dsn, threads):
 201     """ Copy data into the word and placex table.
 202     """
 203     sel = selectors.DefaultSelector()
 204     # Then copy data from place to placex in <threads - 1> chunks.
 205     place_threads = max(1, threads - 1)
 206     for imod in range(place_threads):
 207         conn = DBConnection(dsn)
 208         conn.connect()
 209         conn.perform(
 210             pysql.SQL("""INSERT INTO placex ({columns})
 211                            SELECT {columns} FROM place
 212                            WHERE osm_id % {total} = {mod}
 213                              AND NOT (class='place' and (type='houses' or type='postcode'))
 214                              AND ST_IsValid(geometry)
 215                       """).format(columns=_COPY_COLUMNS,
 216                                   total=pysql.Literal(place_threads),
 217                                   mod=pysql.Literal(imod)))
 218         sel.register(conn, selectors.EVENT_READ, conn)
 219
 220     # Address interpolations go into another table.
 221     conn = DBConnection(dsn)
 222     conn.connect()
 223     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
 224                       SELECT osm_id, address, geometry FROM place
 225                       WHERE class='place' and type='houses' and osm_type='W'
 226                             and ST_GeometryType(geometry) = 'ST_LineString'
 227                  """)
 228     sel.register(conn, selectors.EVENT_READ, conn)
 229
 230     # Now wait for all of them to finish.
 231     todo = place_threads + 1
 232     while todo > 0:
 233         for key, _ in sel.select(1):
 234             conn = key.data
 235             sel.unregister(conn)
 236             conn.wait()
 237             conn.close()
 238             todo -= 1
 239         print('.', end='', flush=True)
 240     print('\n')
 241
 242     with connect(dsn) as conn:
 243         with conn.cursor() as cur:
 244             cur.execute('ANALYSE')
 245
 246
 247 def create_search_indices(conn, config, drop=False):
 248     """ Create tables that have explicit partitioning.
 249     """
 250
 251     # If index creation failed and left an index invalid, they need to be
 252     # cleaned out first, so that the script recreates them.
 253     with conn.cursor() as cur:
 254         cur.execute("""SELECT relname FROM pg_class, pg_index
 255                        WHERE pg_index.indisvalid = false
 256                              AND pg_index.indexrelid = pg_class.oid""")
 257         bad_indices = [row[0] for row in list(cur)]
 258         for idx in bad_indices:
 259             LOG.info("Drop invalid index %s.", idx)
 260             cur.execute('DROP INDEX "{}"'.format(idx))
 261     conn.commit()
 262
 263     sql = SQLPreprocessor(conn, config)
 264
 265     sql.run_sql_file(conn, 'indices.sql', drop=drop)
 266
 267
 268 def create_country_names(conn, tokenizer, languages=None):
 269     """ Add default country names to search index. `languages` is a comma-
 270         separated list of language codes as used in OSM. If `languages` is not
 271         empty then only name translations for the given languages are added
 272         to the index.
 273     """
 274     if languages:
 275         languages = languages.split(',')
 276
 277     def _include_key(key):
 278         return key == 'name' or \
 279                (key.startswith('name:') and (not languages or key[5:] in languages))
 280
 281     with conn.cursor() as cur:
 282         psycopg2.extras.register_hstore(cur)
 283         cur.execute("""SELECT country_code, name FROM country_name
 284                        WHERE country_code is not null""")
 285
 286         with tokenizer.name_analyzer() as analyzer:
 287             for code, name in cur:
 288                 names = {'countrycode': code}
 289                 if code == 'gb':
 290                     names['short_name'] = 'UK'
 291                 if code == 'us':
 292                     names['short_name'] = 'United States'
 293
 294                 # country names (only in languages as provided)
 295                 if name:
 296                     names.update(((k, v) for k, v in name.items() if _include_key(k)))
 297
 298                 analyzer.add_country_names(code, names)
 299
 300     conn.commit()