]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/tools/database_import.py
prot load-data function to python
[nominatim.git] / 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')