pylint: disable check too-few-public-methods

[nominatim.git] / nominatim / indexer / indexer.py

diff --git a/nominatim/indexer/indexer.py b/nominatim/indexer/indexer.py
index d997e522490735994080df726e256662a4c0c632..faf91c456b1eb9bd0d6a00d8e1f396f8fd846ab7 100644 (file)
--- a/nominatim/indexer/indexer.py
+++ b/nominatim/indexer/indexer.py
@@ -1,123 +1,79 @@
 """
 Main work horse for indexing (computing addresses) the database.
 """
 """
 Main work horse for indexing (computing addresses) the database.
 """

-# pylint: disable=C0111

 import logging
 import select
 
 import logging
 import select
 

-import psycopg2
-
-from .progress import ProgressLogger
-from ..db.async_connection import DBConnection
+from nominatim.indexer.progress import ProgressLogger
+from nominatim.indexer import runners
+from nominatim.db.async_connection import DBConnection
+from nominatim.db.connection import connect

 
 LOG = logging.getLogger()
 
 
 LOG = logging.getLogger()
 

-class RankRunner:
-    """ Returns SQL commands for indexing one rank within the placex table.
-    """
-
-    def __init__(self, rank):
-        self.rank = rank
-
-    def name(self):
-        return "rank {}".format(self.rank)
-
-    def sql_count_objects(self):
-        return """SELECT count(*) FROM placex
-                  WHERE rank_address = {} and indexed_status > 0
-               """.format(self.rank)
+class WorkerPool:
+    """ A pool of asynchronous database connections.

 
 

-    def sql_get_objects(self):
-        return """SELECT place_id FROM placex
-                  WHERE indexed_status > 0 and rank_address = {}
-                  ORDER BY geometry_sector""".format(self.rank)
+        The pool may be used as a context manager.
+    """
+    REOPEN_CONNECTIONS_AFTER = 100000

 
 

-    @staticmethod
-    def sql_index_place(ids):
-        return "UPDATE placex SET indexed_status = 0 WHERE place_id IN ({})"\
-               .format(','.join((str(i) for i in ids)))
+    def __init__(self, dsn, pool_size):
+        self.threads = [DBConnection(dsn) for _ in range(pool_size)]
+        self.free_workers = self._yield_free_worker()

 
 
 
 

-class InterpolationRunner:
-    """ Returns SQL commands for indexing the address interpolation table
-        location_property_osmline.
-    """
+    def finish_all(self):
+        """ Wait for all connection to finish.
+        """
+        for thread in self.threads:
+            while not thread.is_done():
+                thread.wait()

 
 

-    @staticmethod
-    def name():
-        return "interpolation lines (location_property_osmline)"
-
-    @staticmethod
-    def sql_count_objects():
-        return """SELECT count(*) FROM location_property_osmline
-                  WHERE indexed_status > 0"""
-
-    @staticmethod
-    def sql_get_objects():
-        return """SELECT place_id FROM location_property_osmline
-                  WHERE indexed_status > 0
-                  ORDER BY geometry_sector"""
-
-    @staticmethod
-    def sql_index_place(ids):
-        return """UPDATE location_property_osmline
-                  SET indexed_status = 0 WHERE place_id IN ({})
-               """.format(','.join((str(i) for i in ids)))
-
-class BoundaryRunner:
-    """ Returns SQL commands for indexing the administrative boundaries
-        of a certain rank.
-    """
+        self.free_workers = self._yield_free_worker()

 
 

-    def __init__(self, rank):
-        self.rank = rank
+    def close(self):
+        """ Close all connections and clear the pool.
+        """
+        for thread in self.threads:
+            thread.close()
+        self.threads = []
+        self.free_workers = None

 
 

-    def name(self):
-        return "boundaries rank {}".format(self.rank)

 
 

-    def sql_count_objects(self):
-        return """SELECT count(*) FROM placex
-                  WHERE indexed_status > 0
-                    AND rank_search = {}
-                    AND class = 'boundary' and type = 'administrative'
-               """.format(self.rank)
+    def next_free_worker(self):
+        """ Get the next free connection.
+        """
+        return next(self.free_workers)

 
 

-    def sql_get_objects(self):
-        return """SELECT place_id FROM placex
-                  WHERE indexed_status > 0 and rank_search = {}
-                        and class = 'boundary' and type = 'administrative'
-                  ORDER BY partition, admin_level
-               """.format(self.rank)

 
 

-    @staticmethod
-    def sql_index_place(ids):
-        return "UPDATE placex SET indexed_status = 0 WHERE place_id IN ({})"\
-               .format(','.join((str(i) for i in ids)))
+    def _yield_free_worker(self):
+        ready = self.threads
+        command_stat = 0
+        while True:
+            for thread in ready:
+                if thread.is_done():
+                    command_stat += 1
+                    yield thread

 
 

+            if command_stat > self.REOPEN_CONNECTIONS_AFTER:
+                for thread in self.threads:
+                    while not thread.is_done():
+                        thread.wait()
+                    thread.connect()
+                ready = self.threads
+                command_stat = 0
+            else:
+                _, ready, _ = select.select([], self.threads, [])

 
 

-class PostcodeRunner:
-    """ Provides the SQL commands for indexing the location_postcode table.
-    """

 
 

-    @staticmethod
-    def name():
-        return "postcodes (location_postcode)"
+    def __enter__(self):
+        return self

 
 

-    @staticmethod
-    def sql_count_objects():
-        return 'SELECT count(*) FROM location_postcode WHERE indexed_status > 0'

 
 

-    @staticmethod
-    def sql_get_objects():
-        return """SELECT place_id FROM location_postcode
-                  WHERE indexed_status > 0
-                  ORDER BY country_code, postcode"""
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()

 
 

-    @staticmethod
-    def sql_index_place(ids):
-        return """UPDATE location_postcode SET indexed_status = 0
-                  WHERE place_id IN ({})
-               """.format(','.join((str(i) for i in ids)))

 
 class Indexer:
     """ Main indexing routine.
 
 class Indexer:
     """ Main indexing routine.


@@ -126,23 +82,6 @@ class Indexer:
     def __init__(self, dsn, num_threads):
         self.dsn = dsn
         self.num_threads = num_threads
     def __init__(self, dsn, num_threads):
         self.dsn = dsn
         self.num_threads = num_threads

-        self.conn = None
-        self.threads = []
-
-
-    def _setup_connections(self):
-        self.conn = psycopg2.connect(self.dsn)
-        self.threads = [DBConnection(self.dsn) for _ in range(self.num_threads)]
-
-
-    def _close_connections(self):
-        if self.conn:
-            self.conn.close()
-            self.conn = None
-
-        for thread in self.threads:
-            thread.close()
-        threads = []

 
 
     def index_full(self, analyse=True):
 
 
     def index_full(self, analyse=True):


@@ -151,30 +90,32 @@ class Indexer:
             database will be analysed at the appropriate places to
             ensure that database statistics are updated.
         """
             database will be analysed at the appropriate places to
             ensure that database statistics are updated.
         """

-        conn = psycopg2.connect(self.dsn)
+        with connect(self.dsn) as conn:
+            conn.autocommit = True
+
+            if analyse:
+                def _analyze():
+                    with conn.cursor() as cur:
+                        cur.execute('ANALYZE')
+            else:
+                def _analyze():
+                    pass

 
 

-        try:

             self.index_by_rank(0, 4)
             self.index_by_rank(0, 4)

-            self._analyse_db_if(conn, analyse)
+            _analyze()

 
             self.index_boundaries(0, 30)
 
             self.index_boundaries(0, 30)

-            self._analyse_db_if(conn, analyse)
+            _analyze()

 
             self.index_by_rank(5, 25)
 
             self.index_by_rank(5, 25)

-            self._analyse_db_if(conn, analyse)
+            _analyze()

 
             self.index_by_rank(26, 30)
 
             self.index_by_rank(26, 30)

-            self._analyse_db_if(conn, analyse)
+            _analyze()

 
             self.index_postcodes()
 
             self.index_postcodes()

-            self._analyse_db_if(conn, analyse)
-        finally:
-            conn.close()
+            _analyze()

 
 

-    def _analyse_db_if(self, conn, condition):
-        if condition:
-            with conn.cursor() as cur:
-                cur.execute('ANALYSE')

 
     def index_boundaries(self, minrank, maxrank):
         """ Index only administrative boundaries within the given rank range.
 
     def index_boundaries(self, minrank, maxrank):
         """ Index only administrative boundaries within the given rank range.


@@ -182,13 +123,8 @@ class Indexer:
         LOG.warning("Starting indexing boundaries using %s threads",
                     self.num_threads)
 
         LOG.warning("Starting indexing boundaries using %s threads",
                     self.num_threads)
 

-        self._setup_connections()
-
-        try:
-            for rank in range(max(minrank, 4), min(maxrank, 26)):
-                self.index(BoundaryRunner(rank))
-        finally:
-            self._close_connections()
+        for rank in range(max(minrank, 4), min(maxrank, 26)):
+            self._index(runners.BoundaryRunner(rank))

 
     def index_by_rank(self, minrank, maxrank):
         """ Index all entries of placex in the given rank range (inclusive)
 
     def index_by_rank(self, minrank, maxrank):
         """ Index all entries of placex in the given rank range (inclusive)


@@ -201,20 +137,15 @@ class Indexer:
         LOG.warning("Starting indexing rank (%i to %i) using %i threads",
                     minrank, maxrank, self.num_threads)
 
         LOG.warning("Starting indexing rank (%i to %i) using %i threads",
                     minrank, maxrank, self.num_threads)
 

-        self._setup_connections()
+        for rank in range(max(1, minrank), maxrank):
+            self._index(runners.RankRunner(rank))

 
 

-        try:
-            for rank in range(max(1, minrank), maxrank):
-                self.index(RankRunner(rank))
-
-            if maxrank == 30:
-                self.index(RankRunner(0))
-                self.index(InterpolationRunner(), 20)
-                self.index(RankRunner(30), 20)
-            else:
-                self.index(RankRunner(maxrank))
-        finally:
-            self._close_connections()
+        if maxrank == 30:
+            self._index(runners.RankRunner(0))
+            self._index(runners.InterpolationRunner(), 20)
+            self._index(runners.RankRunner(30), 20)
+        else:
+            self._index(runners.RankRunner(maxrank))

 
 
     def index_postcodes(self):
 
 
     def index_postcodes(self):


@@ -222,89 +153,52 @@ class Indexer:
         """
         LOG.warning("Starting indexing postcodes using %s threads", self.num_threads)
 
         """
         LOG.warning("Starting indexing postcodes using %s threads", self.num_threads)
 

-        self._setup_connections()
+        self._index(runners.PostcodeRunner(), 20)

 
 

-        try:
-            self.index(PostcodeRunner(), 20)
-        finally:
-            self._close_connections()

 
     def update_status_table(self):
         """ Update the status in the status table to 'indexed'.
         """
 
     def update_status_table(self):
         """ Update the status in the status table to 'indexed'.
         """

-        conn = psycopg2.connect(self.dsn)
-
-        try:
+        with connect(self.dsn) as conn:

             with conn.cursor() as cur:
                 cur.execute('UPDATE import_status SET indexed = true')
 
             conn.commit()
             with conn.cursor() as cur:
                 cur.execute('UPDATE import_status SET indexed = true')
 
             conn.commit()

-        finally:
-            conn.close()

 
 

-    def index(self, obj, batch=1):
-        """ Index a single rank or table. `obj` describes the SQL to use
+    def _index(self, runner, batch=1):
+        """ Index a single rank or table. `runner` describes the SQL to use

             for indexing. `batch` describes the number of objects that
             should be processed with a single SQL statement
         """
             for indexing. `batch` describes the number of objects that
             should be processed with a single SQL statement
         """

-        LOG.warning("Starting %s (using batch size %s)", obj.name(), batch)
+        LOG.warning("Starting %s (using batch size %s)", runner.name(), batch)

 
 

-        cur = self.conn.cursor()
-        cur.execute(obj.sql_count_objects())
-
-        total_tuples = cur.fetchone()[0]
-        LOG.debug("Total number of rows: %i", total_tuples)
+        with connect(self.dsn) as conn:
+            with conn.cursor() as cur:
+                total_tuples = cur.scalar(runner.sql_count_objects())
+                LOG.debug("Total number of rows: %i", total_tuples)

 
 

-        cur.close()
+            conn.commit()

 
 

-        progress = ProgressLogger(obj.name(), total_tuples)
+            progress = ProgressLogger(runner.name(), total_tuples)

 
 

-        if total_tuples > 0:
-            cur = self.conn.cursor(name='places')
-            cur.execute(obj.sql_get_objects())
+            if total_tuples > 0:
+                with conn.cursor(name='places') as cur:
+                    cur.execute(runner.sql_get_objects())

 
 

-            next_thread = self.find_free_thread()
-            while True:
-                places = [p[0] for p in cur.fetchmany(batch)]
-                if not places:
-                    break
+                    with WorkerPool(self.dsn, self.num_threads) as pool:
+                        while True:
+                            places = [p[0] for p in cur.fetchmany(batch)]
+                            if not places:
+                                break

 
 

-                LOG.debug("Processing places: %s", str(places))
-                thread = next(next_thread)
+                            LOG.debug("Processing places: %s", str(places))
+                            worker = pool.next_free_worker()

 
 

-                thread.perform(obj.sql_index_place(places))
-                progress.add(len(places))
+                            worker.perform(runner.sql_index_place(places))
+                            progress.add(len(places))

 
 

-            cur.close()
+                        pool.finish_all()

 
 

-            for thread in self.threads:
-                thread.wait()
+                conn.commit()

 
         progress.done()
 
         progress.done()

-
-    def find_free_thread(self):
-        """ Generator that returns the next connection that is free for
-            sending a query.
-        """
-        ready = self.threads
-        command_stat = 0
-
-        while True:
-            for thread in ready:
-                if thread.is_done():
-                    command_stat += 1
-                    yield thread
-
-            # refresh the connections occasionaly to avoid potential
-            # memory leaks in Postgresql.
-            if command_stat > 100000:
-                for thread in self.threads:
-                    while not thread.is_done():
-                        thread.wait()
-                    thread.connect()
-                command_stat = 0
-                ready = self.threads
-            else:
-                ready, _, _ = select.select(self.threads, [], [])
-
-        assert False, "Unreachable code"