add identifier quoting

This commit is contained in:
Jan Dittberner 2015-01-09 22:01:56 +01:00
parent 3fbf645665
commit 6122164990

View file

@ -9,11 +9,46 @@ from celery.utils.log import get_task_logger
from gvapgsql import settings from gvapgsql import settings
from psycopg2 import connect from psycopg2 import connect
from psycopg2.extensions import ISQLQuote
_LOGGER = get_task_logger(__name__) _LOGGER = get_task_logger(__name__)
class Ident(object):
"""
Wrap a PostgreSQL identifier.
Inspired by
http://osdir.com/ml/python-db-psycopg-devel/2009-03/msg00019.html.
Note that this adapter needs to make an SQL SELECT to correctly quote the
identifier (that is then cached) so it can be a good idea to pre-allocate
identifiers by calling getquoted() on the, before entering critical
sections of code.
"""
def __init__(self, ident):
self.ident = ident
self.ident_quoted = None
def prepare(self, conn):
self._conn = conn
def getquoted(self):
if not self.ident_quoted:
with self._conn.cursor() as curs:
curs.execute("SELECT quote_ident(%s)", (self.ident,))
return self.ident_quoted
def __str__(self):
return str(self.ident_quoted)
def __conform__(self, proto):
if proto == ISQLQuote:
return self
def _get_connection(): def _get_connection():
return connect( return connect(
host=settings.GVAPGSQL_DBADMIN_HOST, host=settings.GVAPGSQL_DBADMIN_HOST,
@ -41,7 +76,7 @@ def create_pgsql_user(username, password):
""" """
CREATE USER %(username)s WITH PASSWORD %(password)s CREATE USER %(username)s WITH PASSWORD %(password)s
""", """,
{'username': username, 'password': password} {'username': Ident(username), 'password': password}
) )
@ -63,7 +98,7 @@ def set_pgsql_userpassword(username, password):
""" """
ALTER ROLE %(username)s WITH PASSWORD %(password)s ALTER ROLE %(username)s WITH PASSWORD %(password)s
""", """,
{'username': username, 'password': password} {'username': Ident(username), 'password': password}
) )
@ -83,7 +118,7 @@ def delete_pgsql_user(username):
""" """
DROP ROLE %(username)s DROP ROLE %(username)s
""", """,
{'username': username} {'username': Ident(username)}
) )
@ -104,7 +139,7 @@ def create_pgsql_database(dbname, username):
""" """
CREATE DATABASE %(dbname)s OWNER %(username)s CREATE DATABASE %(dbname)s OWNER %(username)s
""", """,
{'dbname': dbname, 'username': username} {'dbname': Ident(dbname), 'username': Ident(username)}
) )
@ -124,5 +159,5 @@ def delete_pgsql_database(dbname):
""" """
DROP DATABASE %(dbname)s DROP DATABASE %(dbname)s
""", """,
{'dbname': dbname} {'dbname': Ident(dbname)}
) )