Merge remote-tracking branch 'upstream/master'

[nominatim.git] / nominatim / db / connection.py

diff --git a/nominatim/db/connection.py b/nominatim/db/connection.py
index 1eb3599bee64a8cdfbb8c08972138fc6c47e282a..c60bcfddb8af99faf3bc2ba4e7f7a48b7baf731f 100644 (file)
--- a/nominatim/db/connection.py
+++ b/nominatim/db/connection.py
@@ -1,3 +1,9 @@
+# SPDX-License-Identifier: GPL-2.0-only
+#
+# This file is part of Nominatim. (https://nominatim.org)
+#
+# Copyright (C) 2022 by the Nominatim developer community.
+# For a full list of authors see the git log.

 """
 Specialised connection and cursor functions.
 """
 """
 Specialised connection and cursor functions.
 """


@@ -19,13 +25,24 @@ class _Cursor(psycopg2.extras.DictCursor):
         execution functions.
     """
 
         execution functions.
     """
 

-    def execute(self, query, args=None): # pylint: disable=W0221
+    # pylint: disable=arguments-renamed,arguments-differ
+    def execute(self, query, args=None):

         """ Query execution that logs the SQL query when debugging is enabled.
         """
         LOG.debug(self.mogrify(query, args).decode('utf-8'))
 
         super().execute(query, args)
 
         """ Query execution that logs the SQL query when debugging is enabled.
         """
         LOG.debug(self.mogrify(query, args).decode('utf-8'))
 
         super().execute(query, args)
 

+
+    def execute_values(self, sql, argslist, template=None):
+        """ Wrapper for the psycopg2 convenience function to execute
+            SQL for a list of values.
+        """
+        LOG.debug("SQL execute_values(%s, %s)", sql, argslist)
+
+        psycopg2.extras.execute_values(self, sql, argslist, template=template)
+
+

     def scalar(self, sql, args=None):
         """ Execute query that returns a single value. The value is returned.
             If the query yields more than one row, a ValueError is raised.
     def scalar(self, sql, args=None):
         """ Execute query that returns a single value. The value is returned.
             If the query yields more than one row, a ValueError is raised.


@@ -74,6 +91,17 @@ class _Connection(psycopg2.extensions.connection):
             return num == 1
 
 
             return num == 1
 
 

+    def table_has_column(self, table, column):
+        """ Check if the table 'table' exists and has a column with name 'column'.
+        """
+        with self.cursor() as cur:
+            has_column = cur.scalar("""SELECT count(*) FROM information_schema.columns
+                                       WHERE table_name = %s
+                                             and column_name = %s""",
+                                    (table, column))
+            return has_column > 0
+
+

     def index_exists(self, index, table=None):
         """ Check that an index with the given name exists in the database.
             If table is not None then the index must relate to the given
     def index_exists(self, index, table=None):
         """ Check that an index with the given name exists in the database.
             If table is not None then the index must relate to the given


@@ -136,7 +164,7 @@ def connect(dsn):
         ctxmgr.connection = conn
         return ctxmgr
     except psycopg2.OperationalError as err:
         ctxmgr.connection = conn
         return ctxmgr
     except psycopg2.OperationalError as err:

-        raise UsageError("Cannot connect to database: {}".format(err)) from err
+        raise UsageError(f"Cannot connect to database: {err}") from err

 
 
 # Translation from PG connection string parameters to PG environment variables.
 
 
 # Translation from PG connection string parameters to PG environment variables.