diff --git a/NEWS b/NEWS index b269bb3f..fa11f220 100644 --- a/NEWS +++ b/NEWS @@ -7,6 +7,8 @@ What's new in psycopg 2.4.6 - Dropped GIL release during string adaptation around a function call invoking a Python API function, which could cause interpreter crash. Thanks to Manu Cupcic for the report (ticket #110). + - 'register_hstore()', 'register_composite()' work with RealDictConnection + and Cursor. Thanks to Adrian Klaver for the report (ticket #114). - connection.reset() implemented using DISCARD ALL on server versions supporting it. diff --git a/lib/extras.py b/lib/extras.py index 89939728..696a9af6 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -571,6 +571,18 @@ def wait_select(conn): raise OperationalError("bad state from poll: %s" % state) +def _solve_conn_curs(conn_or_curs): + """Return the connection and a DBAPI cursor from a connection or cursor.""" + if hasattr(conn_or_curs, 'execute'): + conn = conn_or_curs.connection + curs = conn.cursor(cursor_factory=_cursor) + else: + conn = conn_or_curs + curs = conn.cursor(cursor_factory=_cursor) + + return conn, curs + + class HstoreAdapter(object): """Adapt a Python dict to the hstore syntax.""" def __init__(self, wrapped): @@ -679,12 +691,7 @@ class HstoreAdapter(object): def get_oids(self, conn_or_curs): """Return the lists of OID of the hstore and hstore[] types. """ - if hasattr(conn_or_curs, 'execute'): - conn = conn_or_curs.connection - curs = conn_or_curs - else: - conn = conn_or_curs - curs = conn_or_curs.cursor() + conn, curs = _solve_conn_curs(conn_or_curs) # Store the transaction status of the connection to revert it after use conn_status = conn.status @@ -890,12 +897,7 @@ class CompositeCaster(object): Raise `ProgrammingError` if the type is not found. """ - if hasattr(conn_or_curs, 'execute'): - conn = conn_or_curs.connection - curs = conn_or_curs - else: - conn = conn_or_curs - curs = conn_or_curs.cursor() + conn, curs = _solve_conn_curs(conn_or_curs) # Store the transaction status of the connection to revert it after use conn_status = conn.status diff --git a/tests/test_types_extras.py b/tests/test_types_extras.py index 36dd77ed..4c63db3f 100755 --- a/tests/test_types_extras.py +++ b/tests/test_types_extras.py @@ -424,6 +424,30 @@ class HstoreTestCase(unittest.TestCase): psycopg2.extensions.string_types.pop(oid) psycopg2.extensions.string_types.pop(aoid) + @skip_if_no_hstore + def test_non_dbapi_connection(self): + from psycopg2.extras import RealDictConnection + from psycopg2.extras import register_hstore + + conn = psycopg2.connect(dsn, connection_factory=RealDictConnection) + try: + register_hstore(conn) + curs = conn.cursor() + curs.execute("select ''::hstore x") + self.assertEqual(curs.fetchone()['x'], {}) + finally: + conn.close() + + conn = psycopg2.connect(dsn, connection_factory=RealDictConnection) + try: + curs = conn.cursor() + register_hstore(curs) + curs.execute("select ''::hstore x") + self.assertEqual(curs.fetchone()['x'], {}) + finally: + conn.close() + + def skip_if_no_composite(f): def skip_if_no_composite_(self): if self.conn.server_version < 80000: @@ -712,6 +736,30 @@ class AdaptTypeTestCase(unittest.TestCase): self.assertEqual(r[0], (2, 'test2')) self.assertEqual(r[1], [(3, 'testc', 2), (4, 'testd', 2)]) + @skip_if_no_hstore + def test_non_dbapi_connection(self): + from psycopg2.extras import RealDictConnection + from psycopg2.extras import register_composite + self._create_type("type_ii", [("a", "integer"), ("b", "integer")]) + + conn = psycopg2.connect(dsn, connection_factory=RealDictConnection) + try: + register_composite('type_ii', conn) + curs = conn.cursor() + curs.execute("select '(1,2)'::type_ii x") + self.assertEqual(curs.fetchone()['x'], (1,2)) + finally: + conn.close() + + conn = psycopg2.connect(dsn, connection_factory=RealDictConnection) + try: + curs = conn.cursor() + register_composite('type_ii', conn) + curs.execute("select '(1,2)'::type_ii x") + self.assertEqual(curs.fetchone()['x'], (1,2)) + finally: + conn.close() + def _create_type(self, name, fields): curs = self.conn.cursor() try: