forked from hans/Nominatim
provide wrapper function for DROP TABLE
Use psycopg2 formatting to ensure correct quoting.
This commit is contained in:
@@ -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()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user