From 91c2ff9296f0b363edbb1880773e749d2905f20d Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 14 Aug 2012 23:29:19 +0100
Subject: [PATCH] Fixed tpc_recover() with RealDictStuff

Same problem and correction of ticket #114.
---
 NEWS                     |  4 ++--
 psycopg/xid_type.c       |  8 ++++++--
 tests/test_connection.py | 18 ++++++++++++++++--
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/NEWS b/NEWS
index fa11f220..d2f8a458 100644
--- a/NEWS
+++ b/NEWS
@@ -7,8 +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).
+  - 'register_hstore()', 'register_composite()', 'tpc_recover()' work with
+    RealDictConnection and Cursor (ticket #114).
   - connection.reset() implemented using DISCARD ALL on server versions
     supporting it.
 
diff --git a/psycopg/xid_type.c b/psycopg/xid_type.c
index b28543ca..b38d9145 100644
--- a/psycopg/xid_type.c
+++ b/psycopg/xid_type.c
@@ -28,6 +28,7 @@
 #include "psycopg/psycopg.h"
 
 #include "psycopg/xid.h"
+#include "psycopg/cursor.h"
 
 
 static const char xid_doc[] =
@@ -660,8 +661,11 @@ xid_recover(PyObject *conn)
     PyObject *tmp;
     Py_ssize_t len, i;
 
-    /* curs = conn.cursor() */
-    if (!(curs = PyObject_CallMethod(conn, "cursor", NULL))) { goto exit; }
+    /* curs = conn.cursor()
+     * (sort of. Use the real cursor in case the connection returns
+     * somenthing non-dbapi -- see ticket #114) */
+    if (!(curs = PyObject_CallFunctionObjArgs(
+        (PyObject *)&cursorType, conn, NULL))) { goto exit; }
 
     /* curs.execute(...) */
     if (!(tmp = PyObject_CallMethod(curs, "execute", "s",
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 07400fa4..e4b7d83c 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -458,8 +458,8 @@ class ConnectionTwoPhaseTests(unittest.TestCase):
         cnn.close()
         return rv
 
-    def connect(self):
-        conn = psycopg2.connect(dsn)
+    def connect(self, **kwargs):
+        conn = psycopg2.connect(dsn, **kwargs)
         self._conns.append(conn)
         return conn
 
@@ -760,6 +760,20 @@ class ConnectionTwoPhaseTests(unittest.TestCase):
         cnn.tpc_prepare()
         self.assertRaises(psycopg2.ProgrammingError, cnn.cancel)
 
+    def test_tpc_recover_non_dbapi_connection(self):
+        from psycopg2.extras import RealDictConnection
+        cnn = self.connect(connection_factory=RealDictConnection)
+        cnn.tpc_begin('dict-connection')
+        cnn.tpc_prepare()
+        cnn.reset()
+
+        xids = cnn.tpc_recover()
+        xid = [ xid for xid in xids if xid.database == dbname ][0]
+        self.assertEqual(None, xid.format_id)
+        self.assertEqual('dict-connection', xid.gtrid)
+        self.assertEqual(None, xid.bqual)
+
+
 from testutils import skip_if_tpc_disabled
 decorate_all_tests(ConnectionTwoPhaseTests, skip_if_tpc_disabled)