nominatim/db/async_connection.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   2 #
   3 # This file is part of Nominatim.
   4 # Copyright (C) 2021 by the Nominatim developer community.
   5 # For a full list of authors see the git log.
   6 """ Database helper functions for the indexer.
   7 """
   8 import logging
   9 import psycopg2
  10 from psycopg2.extras import wait_select
  11
  12 # psycopg2 emits different exceptions pre and post 2.8. Detect if the new error
  13 # module is available and adapt the error handling accordingly.
  14 try:
  15     import psycopg2.errors # pylint: disable=no-name-in-module,import-error
  16     __has_psycopg2_errors__ = True
  17 except ImportError:
  18     __has_psycopg2_errors__ = False
  19
  20 LOG = logging.getLogger()
  21
  22 class DeadlockHandler:
  23     """ Context manager that catches deadlock exceptions and calls
  24         the given handler function. All other exceptions are passed on
  25         normally.
  26     """
  27
  28     def __init__(self, handler):
  29         self.handler = handler
  30
  31     def __enter__(self):
  32         pass
  33
  34     def __exit__(self, exc_type, exc_value, traceback):
  35         if __has_psycopg2_errors__:
  36             if exc_type == psycopg2.errors.DeadlockDetected: # pylint: disable=E1101
  37                 self.handler()
  38                 return True
  39         else:
  40             if exc_type == psycopg2.extensions.TransactionRollbackError:
  41                 if exc_value.pgcode == '40P01':
  42                     self.handler()
  43                     return True
  44         return False
  45
  46
  47 class DBConnection:
  48     """ A single non-blocking database connection.
  49     """
  50
  51     def __init__(self, dsn, cursor_factory=None):
  52         self.current_query = None
  53         self.current_params = None
  54         self.dsn = dsn
  55
  56         self.conn = None
  57         self.cursor = None
  58         self.connect(cursor_factory=cursor_factory)
  59
  60     def close(self):
  61         """ Close all open connections. Does not wait for pending requests.
  62         """
  63         if self.conn is not None:
  64             self.cursor.close()
  65             self.conn.close()
  66
  67         self.conn = None
  68
  69     def connect(self, cursor_factory=None):
  70         """ (Re)connect to the database. Creates an asynchronous connection
  71             with JIT and parallel processing disabled. If a connection was
  72             already open, it is closed and a new connection established.
  73             The caller must ensure that no query is pending before reconnecting.
  74         """
  75         self.close()
  76
  77         # Use a dict to hand in the parameters because async is a reserved
  78         # word in Python3.
  79         self.conn = psycopg2.connect(**{'dsn' : self.dsn, 'async' : True})
  80         self.wait()
  81
  82         self.cursor = self.conn.cursor(cursor_factory=cursor_factory)
  83         # Disable JIT and parallel workers as they are known to cause problems.
  84         # Update pg_settings instead of using SET because it does not yield
  85         # errors on older versions of Postgres where the settings are not
  86         # implemented.
  87         self.perform(
  88             """ UPDATE pg_settings SET setting = -1 WHERE name = 'jit_above_cost';
  89                 UPDATE pg_settings SET setting = 0
  90                    WHERE name = 'max_parallel_workers_per_gather';""")
  91         self.wait()
  92
  93     def _deadlock_handler(self):
  94         LOG.info("Deadlock detected (params = %s), retry.", str(self.current_params))
  95         self.cursor.execute(self.current_query, self.current_params)
  96
  97     def wait(self):
  98         """ Block until any pending operation is done.
  99         """
 100         while True:
 101             with DeadlockHandler(self._deadlock_handler):
 102                 wait_select(self.conn)
 103                 self.current_query = None
 104                 return
 105
 106     def perform(self, sql, args=None):
 107         """ Send SQL query to the server. Returns immediately without
 108             blocking.
 109         """
 110         self.current_query = sql
 111         self.current_params = args
 112         self.cursor.execute(sql, args)
 113
 114     def fileno(self):
 115         """ File descriptor to wait for. (Makes this class select()able.)
 116         """
 117         return self.conn.fileno()
 118
 119     def is_done(self):
 120         """ Check if the connection is available for a new query.
 121
 122             Also checks if the previous query has run into a deadlock.
 123             If so, then the previous query is repeated.
 124         """
 125         if self.current_query is None:
 126             return True
 127
 128         with DeadlockHandler(self._deadlock_handler):
 129             if self.conn.poll() == psycopg2.extensions.POLL_OK:
 130                 self.current_query = None
 131                 return True
 132
 133         return False