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
  12 from nominatim.db.connection import connect, get_pg_env
  13 from nominatim.db import utils as db_utils
  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 setup_database_skeleton(dsn, data_dir, no_partitions, rouser=None):
  23     """ Create a new database for Nominatim and populate it with the
  24         essential extensions and data.
  25     """
  26     LOG.warning('Creating database')
  27     create_db(dsn, rouser)
  28
  29     LOG.warning('Setting up database')
  30     with connect(dsn) as conn:
  31         setup_extensions(conn)
  32
  33     LOG.warning('Loading basic data')
  34     import_base_data(dsn, data_dir, no_partitions)
  35
  36
  37 def create_db(dsn, rouser=None):
  38     """ Create a new database for the given DSN. Fails when the database
  39         already exists or the PostgreSQL version is too old.
  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         postgres_version = conn.server_version_tuple()
  54         if postgres_version < POSTGRESQL_REQUIRED_VERSION:
  55             LOG.fatal('Minimum supported version of Postgresql is %d.%d. '
  56                       'Found version %d.%d.',
  57                       POSTGRESQL_REQUIRED_VERSION[0], POSTGRESQL_REQUIRED_VERSION[1],
  58                       postgres_version[0], postgres_version[1])
  59             raise UsageError('PostgreSQL server is too old.')
  60
  61         if rouser is not None:
  62             with conn.cursor() as cur:
  63                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
  64                                  (rouser, ))
  65                 if cnt == 0:
  66                     LOG.fatal("Web user '%s' does not exists. Create it with:\n"
  67                               "\n      createuser %s", rouser, rouser)
  68                     raise UsageError('Missing read-only user.')
  69
  70
  71
  72 def setup_extensions(conn):
  73     """ Set up all extensions needed for Nominatim. Also checks that the
  74         versions of the extensions are sufficient.
  75     """
  76     with conn.cursor() as cur:
  77         cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
  78         cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
  79     conn.commit()
  80
  81     postgis_version = conn.postgis_version_tuple()
  82     if postgis_version < POSTGIS_REQUIRED_VERSION:
  83         LOG.fatal('Minimum supported version of PostGIS is %d.%d. '
  84                   'Found version %d.%d.',
  85                   POSTGIS_REQUIRED_VERSION[0], POSTGIS_REQUIRED_VERSION[1],
  86                   postgis_version[0], postgis_version[1])
  87         raise UsageError('PostGIS version is too old.')
  88
  89
  90 def import_base_data(dsn, sql_dir, ignore_partitions=False):
  91     """ Create and populate the tables with basic static data that provides
  92         the background for geocoding. Data is assumed to not yet exist.
  93     """
  94     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
  95     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
  96
  97     if ignore_partitions:
  98         with connect(dsn) as conn:
  99             with conn.cursor() as cur:
 100                 cur.execute('UPDATE country_name SET partition = 0')
 101             conn.commit()
 102
 103
 104 def import_osm_data(osm_file, options, drop=False, ignore_errors=False):
 105     """ Import the given OSM file. 'options' contains the list of
 106         default settings for osm2pgsql.
 107     """
 108     options['import_file'] = osm_file
 109     options['append'] = False
 110     options['threads'] = 1
 111
 112     if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
 113         # Make some educated guesses about cache size based on the size
 114         # of the import file and the available memory.
 115         mem = psutil.virtual_memory()
 116         fsize = os.stat(str(osm_file)).st_size
 117         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
 118                                              fsize * 2) / 1024 / 1024) + 1
 119
 120     run_osm2pgsql(options)
 121
 122     with connect(options['dsn']) as conn:
 123         if not ignore_errors:
 124             with conn.cursor() as cur:
 125                 cur.execute('SELECT * FROM place LIMIT 1')
 126                 if cur.rowcount == 0:
 127                     raise UsageError('No data imported by osm2pgsql.')
 128
 129         if drop:
 130             conn.drop_table('planet_osm_nodes')
 131
 132     if drop:
 133         if 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, max_word_frequency=None):
 164     """ Truncate all data tables to prepare for a fresh load.
 165     """
 166     with conn.cursor() as cur:
 167         cur.execute('TRUNCATE word')
 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         if max_word_frequency is not None:
 187             # Used by getorcreate_word_id to ignore frequent partial words.
 188             cur.execute("""CREATE OR REPLACE FUNCTION get_maxwordfreq()
 189                            RETURNS integer AS $$
 190                              SELECT {} as maxwordfreq;
 191                            $$ LANGUAGE SQL IMMUTABLE
 192                         """.format(max_word_frequency))
 193         conn.commit()
 194
 195 _COPY_COLUMNS = 'osm_type, osm_id, class, type, name, admin_level, address, extratags, geometry'
 196
 197 def load_data(dsn, data_dir, threads):
 198     """ Copy data into the word and placex table.
 199     """
 200     # Pre-calculate the most important terms in the word list.
 201     db_utils.execute_file(dsn, data_dir / 'words.sql')
 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("""INSERT INTO placex ({0})
 210                          SELECT {0} FROM place
 211                          WHERE osm_id % {1} = {2}
 212                            AND NOT (class='place' and type='houses')
 213                            AND ST_IsValid(geometry)
 214                      """.format(_COPY_COLUMNS, place_threads, imod))
 215         sel.register(conn, selectors.EVENT_READ, conn)
 216
 217     # Address interpolations go into another table.
 218     conn = DBConnection(dsn)
 219     conn.connect()
 220     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
 221                       SELECT osm_id, address, geometry FROM place
 222                       WHERE class='place' and type='houses' and osm_type='W'
 223                             and ST_GeometryType(geometry) = 'ST_LineString'
 224                  """)
 225     sel.register(conn, selectors.EVENT_READ, conn)
 226
 227     # Now wait for all of them to finish.
 228     todo = place_threads + 1
 229     while todo > 0:
 230         for key, _ in sel.select(1):
 231             conn = key.data
 232             sel.unregister(conn)
 233             conn.wait()
 234             conn.close()
 235             todo -= 1
 236         print('.', end='', flush=True)
 237     print('\n')
 238
 239     with connect(dsn) as conn:
 240         with conn.cursor() as cur:
 241             cur.execute('ANALYSE')
 242
 243
 244 def create_search_indices(conn, config, drop=False):
 245     """ Create tables that have explicit partitioning.
 246     """
 247
 248     # If index creation failed and left an index invalid, they need to be
 249     # cleaned out first, so that the script recreates them.
 250     with conn.cursor() as cur:
 251         cur.execute("""SELECT relname FROM pg_class, pg_index
 252                        WHERE pg_index.indisvalid = false
 253                              AND pg_index.indexrelid = pg_class.oid""")
 254         bad_indices = [row[0] for row in list(cur)]
 255         for idx in bad_indices:
 256             LOG.info("Drop invalid index %s.", idx)
 257             cur.execute('DROP INDEX "{}"'.format(idx))
 258     conn.commit()
 259
 260     sql = SQLPreprocessor(conn, config)
 261
 262     sql.run_sql_file(conn, 'indices.sql', drop=drop)
 263
 264 def create_country_names(conn, config):
 265     """ Create search index for default country names.
 266     """
 267
 268     with conn.cursor() as cur:
 269         cur.execute("""SELECT getorcreate_country(make_standard_name('uk'), 'gb')""")
 270         cur.execute("""SELECT getorcreate_country(make_standard_name('united states'), 'us')""")
 271         cur.execute("""SELECT COUNT(*) FROM
 272                        (SELECT getorcreate_country(make_standard_name(country_code),
 273                        country_code) FROM country_name WHERE country_code is not null) AS x""")
 274         cur.execute("""SELECT COUNT(*) FROM
 275                        (SELECT getorcreate_country(make_standard_name(name->'name'), country_code)
 276                        FROM country_name WHERE name ? 'name') AS x""")
 277         sql_statement = """SELECT COUNT(*) FROM (SELECT getorcreate_country(make_standard_name(v),
 278                            country_code) FROM (SELECT country_code, skeys(name)
 279                            AS k, svals(name) AS v FROM country_name) x WHERE k"""
 280
 281         languages = config.LANGUAGES
 282
 283         if languages:
 284             sql_statement = "{} IN (".format(sql_statement)
 285             delim = ''
 286             for language in languages.split(','):
 287                 sql_statement = "{}{}'name:{}'".format(sql_statement, delim, language)
 288                 delim = ', '
 289             sql_statement = '{})'.format(sql_statement)
 290         else:
 291             sql_statement = "{} LIKE 'name:%'".format(sql_statement)
 292         sql_statement = "{}) v".format(sql_statement)
 293         cur.execute(sql_statement)
 294     conn.commit()