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.info("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.info("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.info("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         for r in cur:
 204             sector = r[0]
 205
 206             # Should we do the remaining ones together?
 207             do_all = done_tuples > min_grouped_tuples
 208
 209             pcur = self.conn.cursor(name='places')
 210
 211             if do_all:
 212                 pcur.execute(obj.sql_nosector_places())
 213             else:
 214                 pcur.execute(sector_sql, (sector, ))
 215
 216             for place in pcur:
 217                 place_id = place[0]
 218                 log.debug("Processing place {}".format(place_id))
 219                 thread = next(next_thread)
 220
 221                 thread.perform(index_sql, (place_id,))
 222                 done_tuples += 1
 223
 224             pcur.close()
 225
 226             if do_all:
 227                 break
 228
 229         cur.close()
 230
 231         for t in self.threads:
 232             t.wait()
 233
 234         rank_end_time = datetime.now()
 235         diff_seconds = (rank_end_time-rank_start_time).total_seconds()
 236
 237         log.info("Done {}/{} in {} @ {} per second - FINISHED {}\n".format(
 238                  done_tuples, total_tuples, int(diff_seconds),
 239                  done_tuples/diff_seconds, obj.name()))
 240
 241     def find_free_thread(self):
 242         """ Generator that returns the next connection that is free for
 243             sending a query.
 244         """
 245         ready = self.threads
 246
 247         while True:
 248             for thread in ready:
 249                 if thread.is_done():
 250                     yield thread
 251
 252             ready, _, _ = select.select(self.threads, [], [])
 253
 254         assert(False, "Unreachable code")
 255
 256
 257 def nominatim_arg_parser():
 258     """ Setup the command-line parser for the tool.
 259     """
 260     def h(s):
 261         return re.sub("\s\s+" , " ", s)
 262
 263     p = ArgumentParser(description="Indexing tool for Nominatim.",
 264                        formatter_class=RawDescriptionHelpFormatter)
 265
 266     p.add_argument('-d', '--database',
 267                    dest='dbname', action='store', default='nominatim',
 268                    help='Name of the PostgreSQL database to connect to.')
 269     p.add_argument('-U', '--username',
 270                    dest='user', action='store',
 271                    help='PostgreSQL user name.')
 272     p.add_argument('-W', '--password',
 273                    dest='password_prompt', action='store_true',
 274                    help='Force password prompt.')
 275     p.add_argument('-H', '--host',
 276                    dest='host', action='store',
 277                    help='PostgreSQL server hostname or socket location.')
 278     p.add_argument('-P', '--port',
 279                    dest='port', action='store',
 280                    help='PostgreSQL server port')
 281     p.add_argument('-r', '--minrank',
 282                    dest='minrank', type=int, metavar='RANK', default=0,
 283                    help='Minimum/starting rank.')
 284     p.add_argument('-R', '--maxrank',
 285                    dest='maxrank', type=int, metavar='RANK', default=30,
 286                    help='Maximum/finishing rank.')
 287     p.add_argument('-t', '--threads',
 288                    dest='threads', type=int, metavar='NUM', default=1,
 289                    help='Number of threads to create for indexing.')
 290     p.add_argument('-v', '--verbose',
 291                    dest='loglevel', action='count', default=0,
 292                    help='Increase verbosity')
 293
 294     return p
 295
 296 if __name__ == '__main__':
 297     logging.basicConfig(stream=sys.stderr, format='%(levelname)s: %(message)s')
 298
 299     options = nominatim_arg_parser().parse_args(sys.argv[1:])
 300
 301     log.setLevel(max(3 - options.loglevel, 0) * 10)
 302
 303     options.password = None
 304     if options.password_prompt:
 305         password = getpass.getpass("Database password: ")
 306         options.password = password
 307
 308     Indexer(options).run()