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