provide wrapper function for DROP TABLE

Use psycopg2 formatting to ensure correct quoting.
This commit is contained in:
Sarah Hoffmann
2021-07-12 20:32:46 +02:00
parent cf98cff2a1
commit 06602b4ec0
5 changed files with 26 additions and 16 deletions

View File

@@ -8,6 +8,7 @@ import os
import psycopg2
import psycopg2.extensions
import psycopg2.extras
from psycopg2 import sql as pysql
from nominatim.errors import UsageError
@@ -37,6 +38,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 +92,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()