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 5aa05ced1ceb22cfb0c7e5ab218281ade9d0d18f..1c1152079af9e8437c4f685f1d52bd4afcd40bd3 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.
 """
@@ -8,8 +14,9 @@ import os
 import psycopg2
 import psycopg2.extensions
 import psycopg2.extras
+from psycopg2 import sql as pysql
 
-from ..errors import UsageError
+from nominatim.errors import UsageError
 
 LOG = logging.getLogger()
 
@@ -25,6 +32,16 @@ class _Cursor(psycopg2.extras.DictCursor):
 
         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.
@@ -37,6 +54,22 @@ class _Cursor(psycopg2.extras.DictCursor):
         return self.fetchone()[0]
 
 
+    def drop_table(self, name, if_exists=True, cascade=False):
+        """ Drop the table with the given name.
+            Set `if_exists` to False if a non-existant table should raise
+            an exception instead of just being ignored. If 'cascade' is set
+            to True then all dependent tables are deleted as well.
+        """
+        sql = 'DROP TABLE '
+        if if_exists:
+            sql += 'IF EXISTS '
+        sql += '{}'
+        if cascade:
+            sql += ' CASCADE'
+
+        self.execute(pysql.SQL(sql).format(pysql.Identifier(name)))
+
+
 class _Connection(psycopg2.extensions.connection):
     """ A connection that provides the specialised cursor by default and
         adds convenience functions for administrating the database.
@@ -75,14 +108,13 @@ class _Connection(psycopg2.extensions.connection):
         return True
 
 
-    def drop_table(self, name, if_exists=True):
+    def drop_table(self, name, if_exists=True, cascade=False):
         """ Drop the table with the given name.
             Set `if_exists` to False if a non-existant table should raise
             an exception instead of just being ignored.
         """
         with self.cursor() as cur:
-            cur.execute("""DROP TABLE {} "{}"
-                        """.format('IF EXISTS' if if_exists else '', name))
+            cur.drop_table(name, if_exists, cascade)
         self.commit()