X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/fbbdd31399da42b94188d9d4aa4f084efd4876a4..56d24085f9b5b2e1d0d0fd5ff8af0c905765ff04:/nominatim/tools/database_import.py?ds=sidebyside diff --git a/nominatim/tools/database_import.py b/nominatim/tools/database_import.py index 400ce7a5..0dd93490 100644 --- a/nominatim/tools/database_import.py +++ b/nominatim/tools/database_import.py @@ -8,6 +8,8 @@ import subprocess from pathlib import Path import psutil +import psycopg2.extras +from psycopg2 import sql as pysql from nominatim.db.connection import connect, get_pg_env from nominatim.db import utils as db_utils @@ -101,11 +103,11 @@ def import_base_data(dsn, sql_dir, ignore_partitions=False): conn.commit() -def import_osm_data(osm_file, options, drop=False, ignore_errors=False): - """ Import the given OSM file. 'options' contains the list of +def import_osm_data(osm_files, options, drop=False, ignore_errors=False): + """ Import the given OSM files. 'options' contains the list of default settings for osm2pgsql. """ - options['import_file'] = osm_file + options['import_file'] = osm_files options['append'] = False options['threads'] = 1 @@ -113,7 +115,12 @@ def import_osm_data(osm_file, options, drop=False, ignore_errors=False): # Make some educated guesses about cache size based on the size # of the import file and the available memory. mem = psutil.virtual_memory() - fsize = os.stat(str(osm_file)).st_size + fsize = 0 + if isinstance(osm_files, list): + for fname in osm_files: + fsize += os.stat(str(fname)).st_size + else: + fsize = os.stat(str(osm_files)).st_size options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75, fsize * 2) / 1024 / 1024) + 1 @@ -129,9 +136,8 @@ def import_osm_data(osm_file, options, drop=False, ignore_errors=False): if drop: conn.drop_table('planet_osm_nodes') - if drop: - if options['flatnode_file']: - Path(options['flatnode_file']).unlink() + if drop and options['flatnode_file']: + Path(options['flatnode_file']).unlink() def create_tables(conn, config, reverse_only=False): @@ -184,7 +190,12 @@ def truncate_data_tables(conn): conn.commit() -_COPY_COLUMNS = 'osm_type, osm_id, class, type, name, admin_level, address, extratags, geometry' + +_COPY_COLUMNS = pysql.SQL(',').join(map(pysql.Identifier, + ('osm_type', 'osm_id', 'class', 'type', + 'name', 'admin_level', 'address', + 'extratags', 'geometry'))) + def load_data(dsn, threads): """ Copy data into the word and placex table. @@ -195,12 +206,15 @@ def load_data(dsn, threads): for imod in range(place_threads): conn = DBConnection(dsn) conn.connect() - conn.perform("""INSERT INTO placex ({0}) - SELECT {0} FROM place - WHERE osm_id % {1} = {2} - AND NOT (class='place' and type='houses') - AND ST_IsValid(geometry) - """.format(_COPY_COLUMNS, place_threads, imod)) + conn.perform( + pysql.SQL("""INSERT INTO placex ({columns}) + SELECT {columns} FROM place + WHERE osm_id % {total} = {mod} + AND NOT (class='place' and (type='houses' or type='postcode')) + AND ST_IsValid(geometry) + """).format(columns=_COPY_COLUMNS, + total=pysql.Literal(place_threads), + mod=pysql.Literal(imod))) sel.register(conn, selectors.EVENT_READ, conn) # Address interpolations go into another table. @@ -250,34 +264,37 @@ def create_search_indices(conn, config, drop=False): sql.run_sql_file(conn, 'indices.sql', drop=drop) -def create_country_names(conn, config): - """ Create search index for default country names. + +def create_country_names(conn, tokenizer, languages=None): + """ Add default country names to search index. `languages` is a comma- + separated list of language codes as used in OSM. If `languages` is not + empty then only name translations for the given languages are added + to the index. """ + if languages: + languages = languages.split(',') + + def _include_key(key): + return key == 'name' or \ + (key.startswith('name:') and (not languages or key[5:] in languages)) with conn.cursor() as cur: - cur.execute("""SELECT getorcreate_country(make_standard_name('uk'), 'gb')""") - cur.execute("""SELECT getorcreate_country(make_standard_name('united states'), 'us')""") - cur.execute("""SELECT COUNT(*) FROM - (SELECT getorcreate_country(make_standard_name(country_code), - country_code) FROM country_name WHERE country_code is not null) AS x""") - cur.execute("""SELECT COUNT(*) FROM - (SELECT getorcreate_country(make_standard_name(name->'name'), country_code) - FROM country_name WHERE name ? 'name') AS x""") - sql_statement = """SELECT COUNT(*) FROM (SELECT getorcreate_country(make_standard_name(v), - country_code) FROM (SELECT country_code, skeys(name) - AS k, svals(name) AS v FROM country_name) x WHERE k""" - - languages = config.LANGUAGES - - if languages: - sql_statement = "{} IN (".format(sql_statement) - delim = '' - for language in languages.split(','): - sql_statement = "{}{}'name:{}'".format(sql_statement, delim, language) - delim = ', ' - sql_statement = '{})'.format(sql_statement) - else: - sql_statement = "{} LIKE 'name:%'".format(sql_statement) - sql_statement = "{}) v".format(sql_statement) - cur.execute(sql_statement) + psycopg2.extras.register_hstore(cur) + cur.execute("""SELECT country_code, name FROM country_name + WHERE country_code is not null""") + + with tokenizer.name_analyzer() as analyzer: + for code, name in cur: + names = {'countrycode': code} + if code == 'gb': + names['short_name'] = 'UK' + if code == 'us': + names['short_name'] = 'United States' + + # country names (only in languages as provided) + if name: + names.update(((k, v) for k, v in name.items() if _include_key(k))) + + analyzer.add_country_names(code, names) + conn.commit()