Proper unicode handling in quote_ident.

2 files changed, 41 insertions, 10 deletions

diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 9906b7b..cf70a4a 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -166,17 +166,25 @@ exit:
 }
 
 
-#define psyco_quote_ident_doc "quote_ident(str, conn_or_curs) -> str"
+#define psyco_quote_ident_doc \
+"quote_ident(str, conn_or_curs) -> str -- wrapper around PQescapeIdentifier\n\n" \
+":Parameters:\n" \
+"  * `str`: A bytes or unicode object\n" \
+"  * `conn_or_curs`: A connection or cursor, required"
 
 static PyObject *
-psyco_quote_ident(PyObject *self, PyObject *args)
+psyco_quote_ident(PyObject *self, PyObject *args, PyObject *kwargs)
 {
-    const char *str = NULL;
-    char *quoted;
-    PyObject *obj, *result;
+#if PG_VERSION_NUM >= 90000
+    PyObject *ident = NULL, *obj = NULL, *result = NULL;
     connectionObject *conn;
+    const char *str;
+    char *quoted = NULL;
 
-    if (!PyArg_ParseTuple(args, "sO", &str, &obj)) return NULL;
+    static char *kwlist[] = {"ident", "scope", NULL};
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO", kwlist, &ident, &obj)) {
+        return NULL;
+    }
 
     if (PyObject_TypeCheck(obj, &cursorType)) {
         conn = ((cursorObject*)obj)->conn;
@@ -190,15 +198,27 @@ psyco_quote_ident(PyObject *self, PyObject *args)
         return NULL;
     }
 
+    Py_INCREF(ident); /* for ensure_bytes */
+    if (!(ident = psycopg_ensure_bytes(ident))) { goto exit; }
+
+    str = Bytes_AS_STRING(ident);
+
     quoted = PQescapeIdentifier(conn->pgconn, str, strlen(str));
     if (!quoted) {
         PyErr_NoMemory();
-        return NULL;
+        goto exit;
     }
     result = conn_text_from_chars(conn, quoted);
+
+exit:
     PQfreemem(quoted);
+    Py_XDECREF(ident);
 
     return result;
+#else
+    PyErr_SetString(NotSupportedError, "PQescapeIdentifier not available in libpq < 9.0");
+    return NULL;
+#endif
 }
 
 /** type registration **/
@@ -802,10 +822,10 @@ static PyMethodDef psycopgMethods[] = {
      METH_VARARGS|METH_KEYWORDS, psyco_connect_doc},
     {"parse_dsn",  (PyCFunction)psyco_parse_dsn,
      METH_VARARGS|METH_KEYWORDS, psyco_parse_dsn_doc},
+    {"quote_ident", (PyCFunction)psyco_quote_ident,
+     METH_VARARGS|METH_KEYWORDS, psyco_quote_ident_doc},
     {"adapt",  (PyCFunction)psyco_microprotocols_adapt,
      METH_VARARGS, psyco_microprotocols_adapt_doc},
-    {"quote_ident",  (PyCFunction)psyco_quote_ident,
-     METH_VARARGS, psyco_quote_ident_doc},
 
     {"register_type", (PyCFunction)psyco_register_type,
      METH_VARARGS, psyco_register_type_doc},
diff --git a/tests/test_quote.py b/tests/test_quote.py
index a24ab6d..6e94562 100755
--- a/tests/test_quote.py
+++ b/tests/test_quote.py
@@ -23,7 +23,7 @@
 # License for more details.
 
 import sys
-from testutils import unittest, ConnectingTestCase
+from testutils import unittest, ConnectingTestCase, skip_before_libpq
 
 import psycopg2
 import psycopg2.extensions
@@ -166,11 +166,22 @@ class TestQuotedString(ConnectingTestCase):
 
 
 class TestQuotedIdentifier(ConnectingTestCase):
+    @skip_before_libpq(9, 0)
     def test_identifier(self):
         from psycopg2.extensions import quote_ident
         self.assertEqual(quote_ident('blah-blah', self.conn), '"blah-blah"')
         self.assertEqual(quote_ident('quote"inside', self.conn), '"quote""inside"')
 
+    @skip_before_libpq(9, 0)
+    def test_unicode_ident(self):
+        from psycopg2.extensions import quote_ident
+        snowman = u"\u2603"
+        quoted = '"' + snowman + '"'
+        if sys.version_info[0] < 3:
+            self.assertEqual(quote_ident(snowman, self.conn), quoted.encode('utf8'))
+        else:
+            self.assertEqual(quote_ident(snowman, self.conn), quoted)
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)