src/nominatim_db/clicmd/setup.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2024 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Implementation of the 'import' subcommand.
   9 """
  10 from typing import Optional
  11 import argparse
  12 import logging
  13 from pathlib import Path
  14 import asyncio
  15
  16 import psutil
  17
  18 from ..errors import UsageError
  19 from ..config import Configuration
  20 from ..db.connection import connect
  21 from ..db import status, properties
  22 from ..tokenizer.base import AbstractTokenizer
  23 from ..version import NOMINATIM_VERSION
  24 from .args import NominatimArgs
  25
  26 import time
  27
  28 LOG = logging.getLogger()
  29
  30
  31 class SetupAll:
  32     """\
  33     Create a new Nominatim database from an OSM file.
  34
  35     This sub-command sets up a new Nominatim database from scratch starting
  36     with creating a new database in Postgresql. The user running this command
  37     needs superuser rights on the database.
  38     """
  39
  40     def add_args(self, parser: argparse.ArgumentParser) -> None:
  41         group1 = parser.add_argument_group('Required arguments')
  42         group1.add_argument('--osm-file', metavar='FILE', action='append',
  43                             help='OSM file to be imported'
  44                                  ' (repeat for importing multiple files)',
  45                             default=None)
  46         group1.add_argument('--continue', dest='continue_at',
  47                             choices=['import-from-file', 'load-data', 'indexing', 'db-postprocess'],
  48                             help='Continue an import that was interrupted',
  49                             default=None)
  50         group2 = parser.add_argument_group('Optional arguments')
  51         group2.add_argument('--osm2pgsql-cache', metavar='SIZE', type=int,
  52                             help='Size of cache to be used by osm2pgsql (in MB)')
  53         group2.add_argument('--reverse-only', action='store_true',
  54                             help='Do not create tables and indexes for searching')
  55         group2.add_argument('--no-partitions', action='store_true',
  56                             help="Do not partition search indices "
  57                                  "(speeds up import of single country extracts)")
  58         group2.add_argument('--no-updates', action='store_true',
  59                             help="Do not keep tables that are only needed for "
  60                                  "updating the database later")
  61         group2.add_argument('--offline', action='store_true',
  62                             help="Do not attempt to load any additional data from the internet")
  63         group3 = parser.add_argument_group('Expert options')
  64         group3.add_argument('--ignore-errors', action='store_true',
  65                             help='Continue import even when errors in SQL are present')
  66         group3.add_argument('--index-noanalyse', action='store_true',
  67                             help='Do not perform analyse operations during index (expert only)')
  68         group3.add_argument('--prepare-database', action='store_true',
  69                             help='Create the database but do not import any data')
  70
  71     def run(self, args: NominatimArgs) -> int:
  72         if args.osm_file is None and args.continue_at is None and not args.prepare_database:
  73             raise UsageError("No input files (use --osm-file).")
  74
  75         if args.osm_file is not None and args.continue_at not in ('import-from-file', None):
  76             raise UsageError(f"Cannot use --continue {args.continue_at} and --osm-file together.")
  77
  78         if args.continue_at is not None and args.prepare_database:
  79             raise UsageError(
  80                 "Cannot use --continue and --prepare-database together."
  81             )
  82
  83         return asyncio.run(self.async_run(args))
  84
  85     async def async_run(self, args: NominatimArgs) -> int:
  86         from ..data import country_info
  87         from ..tools import database_import, postcodes, freeze
  88         from ..indexer.indexer import Indexer
  89
  90         start_time = time.time()
  91
  92         num_threads = args.threads or psutil.cpu_count() or 1
  93         country_info.setup_country_config(args.config)
  94
  95         if args.prepare_database or args.continue_at is None:
  96             LOG.warning('Creating database')
  97             database_import.setup_database_skeleton(args.config.get_libpq_dsn(),
  98                                                     rouser=args.config.DATABASE_WEBUSER)
  99             if args.prepare_database:
 100                 return 0
 101
 102         if args.continue_at in (None, 'import-from-file'):
 103             self._base_import(args)
 104
 105         if args.continue_at in ('import-from-file', 'load-data', None):
 106             LOG.warning('Initialise tables')
 107             with connect(args.config.get_libpq_dsn()) as conn:
 108                 database_import.truncate_data_tables(conn)
 109
 110             LOG.warning('Load data into placex table')
 111             await database_import.load_data(args.config.get_libpq_dsn(), num_threads)
 112
 113         LOG.warning("Setting up tokenizer")
 114         tokenizer = self._get_tokenizer(args.continue_at, args.config)
 115
 116         if args.continue_at in ('import-from-file', 'load-data', None):
 117             LOG.warning('Calculate postcodes')
 118             postcodes.update_postcodes(args.config.get_libpq_dsn(),
 119                                        args.project_dir, tokenizer)
 120
 121         if args.continue_at in ('import-from-file', 'load-data', 'indexing', None):
 122             LOG.warning('Indexing places')
 123             indexer = Indexer(args.config.get_libpq_dsn(), tokenizer, num_threads)
 124             await indexer.index_full(analyse=not args.index_noanalyse)
 125
 126         LOG.warning('Post-process tables')
 127         with connect(args.config.get_libpq_dsn()) as conn:
 128             conn.autocommit = True
 129             await database_import.create_search_indices(conn, args.config,
 130                                                         drop=args.no_updates,
 131                                                         threads=num_threads)
 132             LOG.warning('Create search index for default country names.')
 133             conn.autocommit = False
 134             country_info.create_country_names(conn, tokenizer,
 135                                               args.config.get_str_list('LANGUAGES'))
 136             if args.no_updates:
 137                 conn.autocommit = True
 138                 freeze.drop_update_tables(conn)
 139         tokenizer.finalize_import(args.config)
 140
 141         LOG.warning('Recompute word counts')
 142         tokenizer.update_statistics(args.config, threads=num_threads)
 143
 144         end_time = time.time()
 145         elapsed = end_time - start_time
 146         LOG.warning(f'Import completed successfully in {elapsed:.2f} seconds.')
 147
 148         self._finalize_database(args.config.get_libpq_dsn(), args.offline)
 149
 150         return 0
 151
 152     def _base_import(self, args: NominatimArgs) -> None:
 153         from ..tools import database_import, refresh
 154         from ..data import country_info
 155
 156         files = args.get_osm_file_list()
 157         if not files:
 158             raise UsageError("No input files (use --osm-file).")
 159
 160         if args.continue_at in ('import-from-file', None):
 161             # Check if the correct plugins are installed
 162             database_import.check_existing_database_plugins(args.config.get_libpq_dsn())
 163             LOG.warning('Setting up country tables')
 164             country_info.setup_country_tables(args.config.get_libpq_dsn(),
 165                                               args.config.lib_dir.data,
 166                                               args.no_partitions)
 167
 168             LOG.warning('Importing OSM data file')
 169             database_import.import_osm_data(files,
 170                                             args.osm2pgsql_options(0, 1),
 171                                             drop=args.no_updates,
 172                                             ignore_errors=args.ignore_errors)
 173
 174             LOG.warning('Importing wikipedia importance data')
 175             data_path = Path(args.config.WIKIPEDIA_DATA_PATH or args.project_dir)
 176             if refresh.import_wikipedia_articles(args.config.get_libpq_dsn(),
 177                                                  data_path) > 0:
 178                 LOG.error('Wikipedia importance dump file not found. '
 179                           'Calculating importance values of locations will not '
 180                           'use Wikipedia importance data.')
 181
 182             LOG.warning('Importing secondary importance raster data')
 183             if refresh.import_secondary_importance(args.config.get_libpq_dsn(),
 184                                                    args.project_dir) != 0:
 185                 LOG.error('Secondary importance file not imported. '
 186                           'Falling back to default ranking.')
 187
 188             self._setup_tables(args.config, args.reverse_only)
 189
 190     def _setup_tables(self, config: Configuration, reverse_only: bool) -> None:
 191         """ Set up the basic database layout: tables, indexes and functions.
 192         """
 193         from ..tools import database_import, refresh
 194
 195         with connect(config.get_libpq_dsn()) as conn:
 196             conn.autocommit = True
 197             LOG.warning('Create functions (1st pass)')
 198             refresh.create_functions(conn, config, False, False)
 199             LOG.warning('Create tables')
 200             database_import.create_tables(conn, config, reverse_only=reverse_only)
 201             refresh.load_address_levels_from_config(conn, config)
 202             LOG.warning('Create functions (2nd pass)')
 203             refresh.create_functions(conn, config, False, False)
 204             LOG.warning('Create table triggers')
 205             database_import.create_table_triggers(conn, config)
 206             LOG.warning('Create partition tables')
 207             database_import.create_partition_tables(conn, config)
 208             LOG.warning('Create functions (3rd pass)')
 209             refresh.create_functions(conn, config, False, False)
 210
 211     def _get_tokenizer(self, continue_at: Optional[str],
 212                        config: Configuration) -> AbstractTokenizer:
 213         """ Set up a new tokenizer or load an already initialised one.
 214         """
 215         from ..tokenizer import factory as tokenizer_factory
 216
 217         if continue_at in ('import-from-file', 'load-data', None):
 218             # (re)initialise the tokenizer data
 219             return tokenizer_factory.create_tokenizer(config)
 220
 221         # just load the tokenizer
 222         return tokenizer_factory.get_tokenizer_for_db(config)
 223
 224     def _finalize_database(self, dsn: str, offline: bool) -> None:
 225         """ Determine the database date and set the status accordingly.
 226         """
 227         with connect(dsn) as conn:
 228             properties.set_property(conn, 'database_version', str(NOMINATIM_VERSION))
 229
 230             try:
 231                 dbdate = status.compute_database_date(conn, offline)
 232                 status.set_status(conn, dbdate)
 233                 LOG.info('Database is at %s.', dbdate)
 234             except Exception as exc:
 235                 LOG.error('Cannot determine date of database: %s', exc)