nominatim/nominatim.py

   1 #! /usr/bin/env python
   2 #-----------------------------------------------------------------------------
   3 # nominatim - [description]
   4 #-----------------------------------------------------------------------------
   5 #
   6 # Indexing tool for the Nominatim database.
   7 #
   8 # Based on C version by Brian Quinion
   9 #
  10 # This program is free software; you can redistribute it and/or
  11 # modify it under the terms of the GNU General Public License
  12 # as published by the Free Software Foundation; either version 2
  13 # of the License, or (at your option) any later version.
  14 #
  15 # This program is distributed in the hope that it will be useful,
  16 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  17 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  18 # GNU General Public License for more details.
  19 #
  20 # You should have received a copy of the GNU General Public License
  21 # along with this program; if not, write to the Free Software
  22 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
  23 #-----------------------------------------------------------------------------
  24
  25 from argparse import ArgumentParser, RawDescriptionHelpFormatter, ArgumentTypeError
  26 import logging
  27 import sys
  28 import re
  29 import getpass
  30 from datetime import datetime
  31 import psycopg2
  32 from psycopg2.extras import wait_select
  33 import select
  34
  35 log = logging.getLogger()
  36
  37 def make_connection(options, asynchronous=False):
  38     return psycopg2.connect(dbname=options.dbname, user=options.user,
  39                             password=options.password, host=options.host,
  40                             port=options.port, async_=asynchronous)
  41
  42
  43 class RankRunner(object):
  44     """ Returns SQL commands for indexing one rank within the placex table.
  45     """
  46
  47     def __init__(self, rank):
  48         self.rank = rank
  49
  50     def name(self):
  51         return "rank {}".format(self.rank)
  52
  53     def sql_index_sectors(self):
  54         return """SELECT geometry_sector, count(*) FROM placex
  55                   WHERE rank_search = {} and indexed_status > 0
  56                   GROUP BY geometry_sector
  57                   ORDER BY geometry_sector""".format(self.rank)
  58
  59     def sql_nosector_places(self):
  60         return """SELECT place_id FROM placex
  61                   WHERE indexed_status > 0 and rank_search = {}
  62                   ORDER BY geometry_sector""".format(self.rank)
  63
  64     def sql_sector_places(self):
  65         return """SELECT place_id FROM placex
  66                   WHERE indexed_status > 0 and rank_search = {}
  67                         and geometry_sector = %s""".format(self.rank)
  68
  69     def sql_index_place(self):
  70         return "UPDATE placex SET indexed_status = 0 WHERE place_id = %s"
  71
  72
  73 class InterpolationRunner(object):
  74     """ Returns SQL commands for indexing the address interpolation table
  75         location_property_osmline.
  76     """
  77
  78     def name(self):
  79         return "interpolation lines (location_property_osmline)"
  80
  81     def sql_index_sectors(self):
  82         return """SELECT geometry_sector, count(*) FROM location_property_osmline
  83                   WHERE indexed_status > 0
  84                   GROUP BY geometry_sector
  85                   ORDER BY geometry_sector"""
  86
  87     def sql_nosector_places(self):
  88         return """SELECT place_id FROM location_property_osmline
  89                   WHERE indexed_status > 0
  90                   ORDER BY geometry_sector"""
  91
  92     def sql_sector_places(self):
  93         return """SELECT place_id FROM location_property_osmline
  94                   WHERE indexed_status > 0 and geometry_sector = %s
  95                   ORDER BY geometry_sector"""
  96
  97     def sql_index_place(self):
  98         return """UPDATE location_property_osmline
  99                   SET indexed_status = 0 WHERE place_id = %s"""
 100
 101
 102 class DBConnection(object):
 103     """ A signle non-blocking database connection.
 104     """
 105
 106     def __init__(self, options):
 107         self.conn = make_connection(options, asynchronous=True)
 108         self.wait()
 109
 110         self.cursor = self.conn.cursor()
 111
 112         self.current_query = None
 113         self.current_params = None
 114
 115     def wait(self):
 116         """ Block until any pending operation is done.
 117         """
 118         wait_select(self.conn)
 119         self.current_query = None
 120
 121     def perform(self, sql, args=None):
 122         """ Send SQL query to the server. Returns immediately without
 123             blocking.
 124         """
 125         self.current_query = sql
 126         self.current_params = args
 127         self.cursor.execute(sql, args)
 128
 129     def fileno(self):
 130         """ File descriptor to wait for. (Makes this class select()able.)
 131         """
 132         return self.conn.fileno()
 133
 134     def is_done(self):
 135         """ Check if the connection is available for a new query.
 136
 137             Also checks if the previous query has run into a deadlock.
 138             If so, then the previous query is repeated.
 139         """
 140         if self.current_query is None:
 141             return True
 142
 143         try:
 144             if self.conn.poll() == psycopg2.extensions.POLL_OK:
 145                 self.current_query = None
 146                 return True
 147         except psycopg2.extensions.TransactionRollbackError as e:
 148             if e.pgcode == '40P01':
 149                 log.debug("Deadlock detected, retry.")
 150                 self.cursor.execute(self.current_query, self.current_params)
 151             else:
 152                 raise
 153
 154         return False
 155
 156
 157 class Indexer(object):
 158     """ Main indexing routine.
 159     """
 160
 161     def __init__(self, options):
 162         self.minrank = max(0, options.minrank)
 163         self.maxrank = min(30, options.maxrank)
 164         self.conn = make_connection(options)
 165         self.threads = [DBConnection(options) for i in range(options.threads)]
 166
 167     def run(self):
 168         """ Run indexing over the entire database.
 169         """
 170         log.warning("Starting indexing rank ({} to {}) using {} threads".format(
 171                  self.minrank, self.maxrank, len(self.threads)))
 172
 173         for rank in range(self.minrank, self.maxrank):
 174             self.index(RankRunner(rank))
 175
 176         if self.maxrank == 30:
 177             self.index(InterpolationRunner())
 178             self.index(RankRunner(30))
 179
 180     def index(self, obj):
 181         """ Index a single rank or table. `obj` describes the SQL to use
 182             for indexing.
 183         """
 184         log.warning("Starting {}".format(obj.name()))
 185
 186         cur = self.conn.cursor(name='main')
 187         cur.execute(obj.sql_index_sectors())
 188
 189         total_tuples = 0
 190         for r in cur:
 191             total_tuples += r[1]
 192         log.debug("Total number of rows; {}".format(total_tuples))
 193
 194         cur.scroll(0, mode='absolute')
 195
 196         next_thread = self.find_free_thread()
 197         done_tuples = 0
 198         rank_start_time = datetime.now()
 199
 200         sector_sql = obj.sql_sector_places()
 201         index_sql = obj.sql_index_place()
 202         min_grouped_tuples = total_tuples - len(self.threads) * 1000
 203
 204         next_info = 100 if log.isEnabledFor(logging.INFO) else total_tuples + 1
 205
 206         for r in cur:
 207             sector = r[0]
 208
 209             # Should we do the remaining ones together?
 210             do_all = done_tuples > min_grouped_tuples
 211
 212             pcur = self.conn.cursor(name='places')
 213
 214             if do_all:
 215                 pcur.execute(obj.sql_nosector_places())
 216             else:
 217                 pcur.execute(sector_sql, (sector, ))
 218
 219             for place in pcur:
 220                 place_id = place[0]
 221                 log.debug("Processing place {}".format(place_id))
 222                 thread = next(next_thread)
 223
 224                 thread.perform(index_sql, (place_id,))
 225                 done_tuples += 1
 226
 227                 if done_tuples >= next_info:
 228                     now = datetime.now()
 229                     done_time = (now - rank_start_time).total_seconds()
 230                     tuples_per_sec = done_tuples / done_time
 231                     log.info("Done {} in {} @ {:.3f} per second - {} ETA (seconds): {:.2f}"
 232                            .format(done_tuples, int(done_time),
 233                                    tuples_per_sec, obj.name(),
 234                                    (total_tuples - done_tuples)/tuples_per_sec))
 235                     next_info += int(tuples_per_sec)
 236
 237             pcur.close()
 238
 239             if do_all:
 240                 break
 241
 242         cur.close()
 243
 244         for t in self.threads:
 245             t.wait()
 246
 247         rank_end_time = datetime.now()
 248         diff_seconds = (rank_end_time-rank_start_time).total_seconds()
 249
 250         log.warning("Done {}/{} in {} @ {:.3f} per second - FINISHED {}\n".format(
 251                  done_tuples, total_tuples, int(diff_seconds),
 252                  done_tuples/diff_seconds, obj.name()))
 253
 254     def find_free_thread(self):
 255         """ Generator that returns the next connection that is free for
 256             sending a query.
 257         """
 258         ready = self.threads
 259
 260         while True:
 261             for thread in ready:
 262                 if thread.is_done():
 263                     yield thread
 264
 265             ready, _, _ = select.select(self.threads, [], [])
 266
 267         assert(False, "Unreachable code")
 268
 269
 270 def nominatim_arg_parser():
 271     """ Setup the command-line parser for the tool.
 272     """
 273     def h(s):
 274         return re.sub("\s\s+" , " ", s)
 275
 276     p = ArgumentParser(description="Indexing tool for Nominatim.",
 277                        formatter_class=RawDescriptionHelpFormatter)
 278
 279     p.add_argument('-d', '--database',
 280                    dest='dbname', action='store', default='nominatim',
 281                    help='Name of the PostgreSQL database to connect to.')
 282     p.add_argument('-U', '--username',
 283                    dest='user', action='store',
 284                    help='PostgreSQL user name.')
 285     p.add_argument('-W', '--password',
 286                    dest='password_prompt', action='store_true',
 287                    help='Force password prompt.')
 288     p.add_argument('-H', '--host',
 289                    dest='host', action='store',
 290                    help='PostgreSQL server hostname or socket location.')
 291     p.add_argument('-P', '--port',
 292                    dest='port', action='store',
 293                    help='PostgreSQL server port')
 294     p.add_argument('-r', '--minrank',
 295                    dest='minrank', type=int, metavar='RANK', default=0,
 296                    help='Minimum/starting rank.')
 297     p.add_argument('-R', '--maxrank',
 298                    dest='maxrank', type=int, metavar='RANK', default=30,
 299                    help='Maximum/finishing rank.')
 300     p.add_argument('-t', '--threads',
 301                    dest='threads', type=int, metavar='NUM', default=1,
 302                    help='Number of threads to create for indexing.')
 303     p.add_argument('-v', '--verbose',
 304                    dest='loglevel', action='count', default=0,
 305                    help='Increase verbosity')
 306
 307     return p
 308
 309 if __name__ == '__main__':
 310     logging.basicConfig(stream=sys.stderr, format='%(levelname)s: %(message)s')
 311
 312     options = nominatim_arg_parser().parse_args(sys.argv[1:])
 313
 314     log.setLevel(max(3 - options.loglevel, 0) * 10)
 315
 316     options.password = None
 317     if options.password_prompt:
 318         password = getpass.getpass("Database password: ")
 319         options.password = password
 320
 321     Indexer(options).run()