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 import shutil
   9 from pathlib import Path
  10
  11 import psutil
  12
  13 from ..db.connection import connect, get_pg_env
  14 from ..db import utils as db_utils
  15 from ..db.async_connection import DBConnection
  16 from .exec_utils import run_osm2pgsql
  17 from ..errors import UsageError
  18 from ..version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
  19
  20 LOG = logging.getLogger()
  21
  22 def create_db(dsn, rouser=None):
  23     """ Create a new database for the given DSN. Fails when the database
  24         already exists or the PostgreSQL version is too old.
  25         Uses `createdb` to create the database.
  26
  27         If 'rouser' is given, then the function also checks that the user
  28         with that given name exists.
  29
  30         Requires superuser rights by the caller.
  31     """
  32     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
  33
  34     if proc.returncode != 0:
  35         raise UsageError('Creating new database failed.')
  36
  37     with connect(dsn) as conn:
  38         postgres_version = conn.server_version_tuple()
  39         if postgres_version < POSTGRESQL_REQUIRED_VERSION:
  40             LOG.fatal('Minimum supported version of Postgresql is %d.%d. '
  41                       'Found version %d.%d.',
  42                       POSTGRESQL_REQUIRED_VERSION[0], POSTGRESQL_REQUIRED_VERSION[1],
  43                       postgres_version[0], postgres_version[1])
  44             raise UsageError('PostgreSQL server is too old.')
  45
  46         if rouser is not None:
  47             with conn.cursor() as cur:
  48                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
  49                                  (rouser, ))
  50                 if cnt == 0:
  51                     LOG.fatal("Web user '%s' does not exists. Create it with:\n"
  52                               "\n      createuser %s", rouser, rouser)
  53                     raise UsageError('Missing read-only user.')
  54
  55
  56
  57 def setup_extensions(conn):
  58     """ Set up all extensions needed for Nominatim. Also checks that the
  59         versions of the extensions are sufficient.
  60     """
  61     with conn.cursor() as cur:
  62         cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
  63         cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
  64     conn.commit()
  65
  66     postgis_version = conn.postgis_version_tuple()
  67     if postgis_version < POSTGIS_REQUIRED_VERSION:
  68         LOG.fatal('Minimum supported version of PostGIS is %d.%d. '
  69                   'Found version %d.%d.',
  70                   POSTGIS_REQUIRED_VERSION[0], POSTGIS_REQUIRED_VERSION[1],
  71                   postgis_version[0], postgis_version[1])
  72         raise UsageError('PostGIS version is too old.')
  73
  74
  75 def install_module(src_dir, project_dir, module_dir):
  76     """ Copy the normalization module from src_dir into the project
  77         directory under the '/module' directory. If 'module_dir' is set, then
  78         use the module from there instead and check that it is accessible
  79         for Postgresql.
  80
  81         The function detects when the installation is run from the
  82         build directory. It doesn't touch the module in that case.
  83     """
  84     if not module_dir:
  85         module_dir = project_dir / 'module'
  86
  87         if not module_dir.exists() or not src_dir.samefile(module_dir):
  88
  89             if not module_dir.exists():
  90                 module_dir.mkdir()
  91
  92             destfile = module_dir / 'nominatim.so'
  93             shutil.copy(str(src_dir / 'nominatim.so'), str(destfile))
  94             destfile.chmod(0o755)
  95
  96             LOG.info('Database module installed at %s', str(destfile))
  97         else:
  98             LOG.info('Running from build directory. Leaving database module as is.')
  99     else:
 100         LOG.info("Using custom path for database module at '%s'", module_dir)
 101
 102     return module_dir
 103
 104
 105 def check_module_dir_path(conn, path):
 106     """ Check that the normalisation module can be found and executed
 107         from the given path.
 108     """
 109     with conn.cursor() as cur:
 110         cur.execute("""CREATE FUNCTION nominatim_test_import_func(text)
 111                        RETURNS text AS '{}/nominatim.so', 'transliteration'
 112                        LANGUAGE c IMMUTABLE STRICT;
 113                        DROP FUNCTION nominatim_test_import_func(text)
 114                     """.format(path))
 115
 116
 117 def import_base_data(dsn, sql_dir, ignore_partitions=False):
 118     """ Create and populate the tables with basic static data that provides
 119         the background for geocoding. Data is assumed to not yet exist.
 120     """
 121     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
 122     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 123
 124     if ignore_partitions:
 125         with connect(dsn) as conn:
 126             with conn.cursor() as cur:
 127                 cur.execute('UPDATE country_name SET partition = 0')
 128             conn.commit()
 129
 130
 131 def import_osm_data(osm_file, options, drop=False):
 132     """ Import the given OSM file. 'options' contains the list of
 133         default settings for osm2pgsql.
 134     """
 135     options['import_file'] = osm_file
 136     options['append'] = False
 137     options['threads'] = 1
 138
 139     if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
 140         # Make some educated guesses about cache size based on the size
 141         # of the import file and the available memory.
 142         mem = psutil.virtual_memory()
 143         fsize = os.stat(str(osm_file)).st_size
 144         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
 145                                              fsize * 2) / 1024 / 1024) + 1
 146
 147     run_osm2pgsql(options)
 148
 149     with connect(options['dsn']) as conn:
 150         with conn.cursor() as cur:
 151             cur.execute('SELECT * FROM place LIMIT 1')
 152             if cur.rowcount == 0:
 153                 raise UsageError('No data imported by osm2pgsql.')
 154
 155         if drop:
 156             conn.drop_table('planet_osm_nodes')
 157
 158     if drop:
 159         if options['flatnode_file']:
 160             Path(options['flatnode_file']).unlink()
 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')
 173         cur.execute('TRUNCATE location_property_tiger')
 174         cur.execute('TRUNCATE location_property_osmline')
 175         cur.execute('TRUNCATE location_postcode')
 176         cur.execute('TRUNCATE search_name')
 177         cur.execute('DROP SEQUENCE 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')