From 61221649905bc1579ad7050c2691284ae88429f7 Mon Sep 17 00:00:00 2001
From: Jan Dittberner <jan@dittberner.info>
Date: Fri, 9 Jan 2015 22:01:56 +0100
Subject: [PATCH] add identifier quoting

---
 gvapgsql/pgsqltasks/tasks.py | 45 ++++++++++++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/gvapgsql/pgsqltasks/tasks.py b/gvapgsql/pgsqltasks/tasks.py
index e28c228..a089245 100644
--- a/gvapgsql/pgsqltasks/tasks.py
+++ b/gvapgsql/pgsqltasks/tasks.py
@@ -9,11 +9,46 @@ from celery.utils.log import get_task_logger
 
 from gvapgsql import settings
 from psycopg2 import connect
+from psycopg2.extensions import ISQLQuote
 
 
 _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():
     return connect(
         host=settings.GVAPGSQL_DBADMIN_HOST,
@@ -41,7 +76,7 @@ def create_pgsql_user(username, password):
                 """
                 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
                 """,
-                {'username': username, 'password': password}
+                {'username': Ident(username), 'password': password}
             )
 
 
@@ -83,7 +118,7 @@ def delete_pgsql_user(username):
                 """
                 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
                 """,
-                {'dbname': dbname, 'username': username}
+                {'dbname': Ident(dbname), 'username': Ident(username)}
             )
 
 
@@ -124,5 +159,5 @@ def delete_pgsql_database(dbname):
                 """
                 DROP DATABASE %(dbname)s
                 """,
-                {'dbname': dbname}
+                {'dbname': Ident(dbname)}
             )