]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/tools/database_import.py
Merge remote-tracking branch 'upstream/master'
[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 from pathlib import Path
9
10 import psutil
11 import psycopg2.extras
12
13 from nominatim.db.connection import connect, get_pg_env
14 from nominatim.db import utils as db_utils
15 from nominatim.db.async_connection import DBConnection
16 from nominatim.db.sql_preprocessor import SQLPreprocessor
17 from nominatim.tools.exec_utils import run_osm2pgsql
18 from nominatim.errors import UsageError
19 from nominatim.version import POSTGRESQL_REQUIRED_VERSION, POSTGIS_REQUIRED_VERSION
20
21 LOG = logging.getLogger()
22
23 def setup_database_skeleton(dsn, data_dir, no_partitions, rouser=None):
24     """ Create a new database for Nominatim and populate it with the
25         essential extensions and data.
26     """
27     LOG.warning('Creating database')
28     create_db(dsn, rouser)
29
30     LOG.warning('Setting up database')
31     with connect(dsn) as conn:
32         setup_extensions(conn)
33
34     LOG.warning('Loading basic data')
35     import_base_data(dsn, data_dir, no_partitions)
36
37
38 def create_db(dsn, rouser=None):
39     """ Create a new database for the given DSN. Fails when the database
40         already exists or the PostgreSQL version is too old.
41         Uses `createdb` to create the database.
42
43         If 'rouser' is given, then the function also checks that the user
44         with that given name exists.
45
46         Requires superuser rights by the caller.
47     """
48     proc = subprocess.run(['createdb'], env=get_pg_env(dsn), check=False)
49
50     if proc.returncode != 0:
51         raise UsageError('Creating new database failed.')
52
53     with connect(dsn) as conn:
54         postgres_version = conn.server_version_tuple()
55         if postgres_version < POSTGRESQL_REQUIRED_VERSION:
56             LOG.fatal('Minimum supported version of Postgresql is %d.%d. '
57                       'Found version %d.%d.',
58                       POSTGRESQL_REQUIRED_VERSION[0], POSTGRESQL_REQUIRED_VERSION[1],
59                       postgres_version[0], postgres_version[1])
60             raise UsageError('PostgreSQL server is too old.')
61
62         if rouser is not None:
63             with conn.cursor() as cur:
64                 cnt = cur.scalar('SELECT count(*) FROM pg_user where usename = %s',
65                                  (rouser, ))
66                 if cnt == 0:
67                     LOG.fatal("Web user '%s' does not exists. Create it with:\n"
68                               "\n      createuser %s", rouser, rouser)
69                     raise UsageError('Missing read-only user.')
70
71
72
73 def setup_extensions(conn):
74     """ Set up all extensions needed for Nominatim. Also checks that the
75         versions of the extensions are sufficient.
76     """
77     with conn.cursor() as cur:
78         cur.execute('CREATE EXTENSION IF NOT EXISTS hstore')
79         cur.execute('CREATE EXTENSION IF NOT EXISTS postgis')
80     conn.commit()
81
82     postgis_version = conn.postgis_version_tuple()
83     if postgis_version < POSTGIS_REQUIRED_VERSION:
84         LOG.fatal('Minimum supported version of PostGIS is %d.%d. '
85                   'Found version %d.%d.',
86                   POSTGIS_REQUIRED_VERSION[0], POSTGIS_REQUIRED_VERSION[1],
87                   postgis_version[0], postgis_version[1])
88         raise UsageError('PostGIS version is too old.')
89
90
91 def import_base_data(dsn, sql_dir, ignore_partitions=False):
92     """ Create and populate the tables with basic static data that provides
93         the background for geocoding. Data is assumed to not yet exist.
94     """
95     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
96     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
97
98     if ignore_partitions:
99         with connect(dsn) as conn:
100             with conn.cursor() as cur:
101                 cur.execute('UPDATE country_name SET partition = 0')
102             conn.commit()
103
104
105 def import_osm_data(osm_file, options, drop=False, ignore_errors=False):
106     """ Import the given OSM file. 'options' contains the list of
107         default settings for osm2pgsql.
108     """
109     options['import_file'] = osm_file
110     options['append'] = False
111     options['threads'] = 1
112
113     if not options['flatnode_file'] and options['osm2pgsql_cache'] == 0:
114         # Make some educated guesses about cache size based on the size
115         # of the import file and the available memory.
116         mem = psutil.virtual_memory()
117         fsize = os.stat(str(osm_file)).st_size
118         options['osm2pgsql_cache'] = int(min((mem.available + mem.cached) * 0.75,
119                                              fsize * 2) / 1024 / 1024) + 1
120
121     run_osm2pgsql(options)
122
123     with connect(options['dsn']) as conn:
124         if not ignore_errors:
125             with conn.cursor() as cur:
126                 cur.execute('SELECT * FROM place LIMIT 1')
127                 if cur.rowcount == 0:
128                     raise UsageError('No data imported by osm2pgsql.')
129
130         if drop:
131             conn.drop_table('planet_osm_nodes')
132
133     if drop:
134         if options['flatnode_file']:
135             Path(options['flatnode_file']).unlink()
136
137
138 def create_tables(conn, config, reverse_only=False):
139     """ Create the set of basic tables.
140         When `reverse_only` is True, then the main table for searching will
141         be skipped and only reverse search is possible.
142     """
143     sql = SQLPreprocessor(conn, config)
144     sql.env.globals['db']['reverse_only'] = reverse_only
145
146     sql.run_sql_file(conn, 'tables.sql')
147
148
149 def create_table_triggers(conn, config):
150     """ Create the triggers for the tables. The trigger functions must already
151         have been imported with refresh.create_functions().
152     """
153     sql = SQLPreprocessor(conn, config)
154     sql.run_sql_file(conn, 'table-triggers.sql')
155
156
157 def create_partition_tables(conn, config):
158     """ Create tables that have explicit partitioning.
159     """
160     sql = SQLPreprocessor(conn, config)
161     sql.run_sql_file(conn, 'partition-tables.src.sql')
162
163
164 def truncate_data_tables(conn):
165     """ Truncate all data tables to prepare for a fresh load.
166     """
167     with conn.cursor() as cur:
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_tiger')
173         cur.execute('TRUNCATE location_property_osmline')
174         cur.execute('TRUNCATE location_postcode')
175         if conn.table_exists('search_name'):
176             cur.execute('TRUNCATE search_name')
177         cur.execute('DROP SEQUENCE IF EXISTS 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     conn.commit()
187
188 _COPY_COLUMNS = 'osm_type, osm_id, class, type, name, admin_level, address, extratags, geometry'
189
190 def load_data(dsn, threads):
191     """ Copy data into the word and placex table.
192     """
193     sel = selectors.DefaultSelector()
194     # Then copy data from place to placex in <threads - 1> chunks.
195     place_threads = max(1, threads - 1)
196     for imod in range(place_threads):
197         conn = DBConnection(dsn)
198         conn.connect()
199         conn.perform("""INSERT INTO placex ({0})
200                          SELECT {0} FROM place
201                          WHERE osm_id % {1} = {2}
202                            AND NOT (class='place' and type='houses')
203                            AND ST_IsValid(geometry)
204                      """.format(_COPY_COLUMNS, place_threads, imod))
205         sel.register(conn, selectors.EVENT_READ, conn)
206
207     # Address interpolations go into another table.
208     conn = DBConnection(dsn)
209     conn.connect()
210     conn.perform("""INSERT INTO location_property_osmline (osm_id, address, linegeo)
211                       SELECT osm_id, address, geometry FROM place
212                       WHERE class='place' and type='houses' and osm_type='W'
213                             and ST_GeometryType(geometry) = 'ST_LineString'
214                  """)
215     sel.register(conn, selectors.EVENT_READ, conn)
216
217     # Now wait for all of them to finish.
218     todo = place_threads + 1
219     while todo > 0:
220         for key, _ in sel.select(1):
221             conn = key.data
222             sel.unregister(conn)
223             conn.wait()
224             conn.close()
225             todo -= 1
226         print('.', end='', flush=True)
227     print('\n')
228
229     with connect(dsn) as conn:
230         with conn.cursor() as cur:
231             cur.execute('ANALYSE')
232
233
234 def create_search_indices(conn, config, drop=False):
235     """ Create tables that have explicit partitioning.
236     """
237
238     # If index creation failed and left an index invalid, they need to be
239     # cleaned out first, so that the script recreates them.
240     with conn.cursor() as cur:
241         cur.execute("""SELECT relname FROM pg_class, pg_index
242                        WHERE pg_index.indisvalid = false
243                              AND pg_index.indexrelid = pg_class.oid""")
244         bad_indices = [row[0] for row in list(cur)]
245         for idx in bad_indices:
246             LOG.info("Drop invalid index %s.", idx)
247             cur.execute('DROP INDEX "{}"'.format(idx))
248     conn.commit()
249
250     sql = SQLPreprocessor(conn, config)
251
252     sql.run_sql_file(conn, 'indices.sql', drop=drop)
253
254 def create_country_names(conn, tokenizer, languages=None):
255     """ Add default country names to search index. `languages` is a comma-
256         separated list of language codes as used in OSM. If `languages` is not
257         empty then only name translations for the given languages are added
258         to the index.
259     """
260     if languages:
261         languages = languages.split(',')
262
263     def _include_key(key):
264         return key == 'name' or \
265                (key.startswith('name:') \
266                 and (not languages or key[5:] in languages))
267
268     with conn.cursor() as cur:
269         psycopg2.extras.register_hstore(cur)
270         cur.execute("""SELECT country_code, name FROM country_name
271                        WHERE country_code is not null""")
272
273         with tokenizer.name_analyzer() as analyzer:
274             for code, name in cur:
275                 names = [code]
276                 if code == 'gb':
277                     names.append('UK')
278                 if code == 'us':
279                     names.append('United States')
280
281                 # country names (only in languages as provided)
282                 if name:
283                     names.extend((v for k, v in name.items() if _include_key(k)))
284
285                 analyzer.add_country_names(code, names)
286
287     conn.commit()