From 9295bce154182863e19342b6a4c2e80a58187120 Mon Sep 17 00:00:00 2001 From: Oleksandr Shulgin Date: Tue, 13 Oct 2015 17:29:55 +0200 Subject: [PATCH 1/3] Add psycopg2.extensions.quote_ident. --- doc/src/extensions.rst | 13 +++++++++++++ lib/extensions.py | 2 +- psycopg/psycopgmodule.c | 38 ++++++++++++++++++++++++++++++++++++++ psycopg/utils.c | 4 ++-- tests/test_quote.py | 7 +++++++ 5 files changed, 61 insertions(+), 3 deletions(-) diff --git a/doc/src/extensions.rst b/doc/src/extensions.rst index 4db76b01..d96cca4f 100644 --- a/doc/src/extensions.rst +++ b/doc/src/extensions.rst @@ -221,6 +221,19 @@ functionalities defined by the |DBAPI|_. .. __: http://www.postgresql.org/docs/current/static/libpq-misc.html#LIBPQ-PQLIBVERSION +.. function:: quote_ident(str, scope) + + Return quoted identifier according to PostgreSQL quoting rules. + + The *scope* must be a `connection` or a `cursor`, the underlying + connection encoding is used for any necessary character conversion. + + Requires libpq >= 9.0. + + .. seealso:: libpq docs for `PQescapeIdentifier()`__ + + .. __: http://www.postgresql.org/docs/current/static/libpq-exec.html#LIBPQ-PQESCAPEIDENTIFIER + .. _sql-adaptation-objects: SQL adaptation protocol objects diff --git a/lib/extensions.py b/lib/extensions.py index d10e8ac6..b40e28b8 100644 --- a/lib/extensions.py +++ b/lib/extensions.py @@ -56,7 +56,7 @@ try: except ImportError: pass -from psycopg2._psycopg import adapt, adapters, encodings, connection, cursor, lobject, Xid, libpq_version, parse_dsn +from psycopg2._psycopg import adapt, adapters, encodings, connection, cursor, lobject, Xid, libpq_version, parse_dsn, quote_ident from psycopg2._psycopg import string_types, binary_types, new_type, new_array_type, register_type from psycopg2._psycopg import ISQLQuote, Notify, Diagnostics, Column diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c index c77dce5b..9906b7be 100644 --- a/psycopg/psycopgmodule.c +++ b/psycopg/psycopgmodule.c @@ -165,6 +165,42 @@ exit: return res; } + +#define psyco_quote_ident_doc "quote_ident(str, conn_or_curs) -> str" + +static PyObject * +psyco_quote_ident(PyObject *self, PyObject *args) +{ + const char *str = NULL; + char *quoted; + PyObject *obj, *result; + connectionObject *conn; + + if (!PyArg_ParseTuple(args, "sO", &str, &obj)) return NULL; + + if (PyObject_TypeCheck(obj, &cursorType)) { + conn = ((cursorObject*)obj)->conn; + } + else if (PyObject_TypeCheck(obj, &connectionType)) { + conn = (connectionObject*)obj; + } + else { + PyErr_SetString(PyExc_TypeError, + "argument 2 must be a connection or a cursor"); + return NULL; + } + + quoted = PQescapeIdentifier(conn->pgconn, str, strlen(str)); + if (!quoted) { + PyErr_NoMemory(); + return NULL; + } + result = conn_text_from_chars(conn, quoted); + PQfreemem(quoted); + + return result; +} + /** type registration **/ #define psyco_register_type_doc \ "register_type(obj, conn_or_curs) -> None -- register obj with psycopg type system\n\n" \ @@ -768,6 +804,8 @@ static PyMethodDef psycopgMethods[] = { METH_VARARGS|METH_KEYWORDS, psyco_parse_dsn_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/psycopg/utils.c b/psycopg/utils.c index 836f6129..ec8e47c8 100644 --- a/psycopg/utils.c +++ b/psycopg/utils.c @@ -87,7 +87,7 @@ psycopg_escape_string(connectionObject *conn, const char *from, Py_ssize_t len, return to; } -/* Escape a string to build a valid PostgreSQL identifier +/* Escape a string to build a valid PostgreSQL identifier. * * Allocate a new buffer on the Python heap containing the new string. * 'len' is optional: if 0 the length is calculated. @@ -96,7 +96,7 @@ psycopg_escape_string(connectionObject *conn, const char *from, Py_ssize_t len, * * WARNING: this function is not so safe to allow untrusted input: it does no * check for multibyte chars. Such a function should be built on - * PQescapeIndentifier, which is only available from PostgreSQL 9.0. + * PQescapeIdentifier, which is only available from PostgreSQL 9.0. */ char * psycopg_escape_identifier_easy(const char *from, Py_ssize_t len) diff --git a/tests/test_quote.py b/tests/test_quote.py index e7b3c316..a24ab6d4 100755 --- a/tests/test_quote.py +++ b/tests/test_quote.py @@ -165,6 +165,13 @@ class TestQuotedString(ConnectingTestCase): self.assertEqual(q.encoding, 'utf_8') +class TestQuotedIdentifier(ConnectingTestCase): + 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"') + + def test_suite(): return unittest.TestLoader().loadTestsFromName(__name__) From 89bb6b0711f8ba59f7b8e81339ddaa53356233a2 Mon Sep 17 00:00:00 2001 From: Oleksandr Shulgin Date: Thu, 15 Oct 2015 11:52:18 +0200 Subject: [PATCH 2/3] Proper unicode handling in quote_ident. --- psycopg/psycopgmodule.c | 38 +++++++++++++++++++++++++++++--------- tests/test_quote.py | 13 ++++++++++++- 2 files changed, 41 insertions(+), 10 deletions(-) diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c index 9906b7be..cf70a4ad 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 a24ab6d4..6e945624 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__) From 109409bc951b7dd3e61712a65289e3458430656a Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Thu, 15 Oct 2015 11:06:44 +0100 Subject: [PATCH 3/3] Mention quote_ident() in NEWS file --- NEWS | 1 + 1 file changed, 1 insertion(+) diff --git a/NEWS b/NEWS index 83ba166e..5200c4dd 100644 --- a/NEWS +++ b/NEWS @@ -14,6 +14,7 @@ New features: - The attributes `~connection.notices` and `~connection.notifies` can be customized replacing them with any object exposing an `!append()` method (:ticket:`#326`). +- Added `~psycopg2.extensions.quote_ident()` function (:ticket:`#359`). What's new in psycopg 2.6.2