move import-data option to native python

[nominatim.git] / nominatim / tools / database_import.py

diff --git a/nominatim/tools/database_import.py b/nominatim/tools/database_import.py
index 89e01f66670fdbbc1567ec23b6c6deccb9375704..00ec95c03ab49553f5eb3bd41b693f25b319efad 100644 (file)
--- a/nominatim/tools/database_import.py
+++ b/nominatim/tools/database_import.py
@@ -2,11 +2,16 @@
 Functions for setting up and importing a new Nominatim database.
 """
 import logging
+import os
 import subprocess
 import shutil
+from pathlib import Path
+
+import psutil
 
 from ..db.connection import connect, get_pg_env
 from ..db import utils as db_utils
+from .exec_utils import run_osm2pgsql
 from ..errors import UsageError
 from ..version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
 
@@ -28,7 +33,7 @@ def create_db(dsn, rouser=None):
         raise UsageError('Creating new database failed.')
 
     with connect(dsn) as conn:
-        postgres_version = conn.server_version_tuple() # pylint: disable=E1101
+        postgres_version = conn.server_version_tuple()
         if postgres_version < POSTGRESQL_REQUIRED_VERSION:
             LOG.fatal('Minimum supported version of Postgresql is %d.%d. '
                       'Found version %d.%d.',
@@ -37,7 +42,7 @@ def create_db(dsn, rouser=None):
             raise UsageError('PostgreSQL server is too old.')
 
         if rouser is not None:
-            with conn.cursor() as cur:  # pylint: disable=E1101
+            with conn.cursor() as cur:
                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
                                  (rouser, ))
                 if cnt == 0:
@@ -109,13 +114,45 @@ def check_module_dir_path(conn, path):
 
 def import_base_data(dsn, sql_dir, ignore_partitions=False):
     """ Create and populate the tables with basic static data that provides
-        the background for geocoding.
+        the background for geocoding. Data is assumed to not yet exist.
     """
     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 
     if ignore_partitions:
         with connect(dsn) as conn:
-            with conn.cursor() as cur:  # pylint: disable=E1101
+            with conn.cursor() as cur:
                 cur.execute('UPDATE country_name SET partition = 0')
-            conn.commit()  # pylint: disable=E1101
+            conn.commit()
+
+
+def import_osm_data(osm_file, options, drop=False):
+    """ Import the given OSM file. 'options' contains the list of
+        default settings for osm2pgsql.
+    """
+    options['import_file'] = osm_file
+    options['append'] = False
+    options['threads'] = 1
+
+    if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
+        # 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
+        options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
+                                             fsize * 2) / 1024 / 1024) + 1
+
+    run_osm2pgsql(options)
+
+    with connect(options['dsn']) as conn:
+        with conn.cursor() as cur:
+            cur.execute('SELECT * FROM place LIMIT 1')
+            if cur.rowcount == 0:
+                raise UsageError('No data imported by osm2pgsql.')
+
+        if drop:
+            conn.drop_table('planet_osm_nodes')
+
+    if drop:
+        if options['flatnode_file']:
+            Path(options['flatnode_file']).unlink()