Invalidate large objects after a two-phase commit operation

2 files changed, 29 insertions, 2 deletions

diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c
index 27e958e..6c7fabf 100644
--- a/psycopg/pqpath.c
+++ b/psycopg/pqpath.c
@@ -610,6 +610,8 @@ pq_tpc_command_locked(connectionObject *conn, const char *cmd, const char *tid,
     Dprintf("_pq_tpc_command: pgconn = %p, command = %s",
             conn->pgconn, cmd);
 
+    conn->mark += 1;
+
     /* convert the xid into the postgres transaction_id and quote it. */
     if (!(etid = psycopg_escape_string((PyObject *)conn, tid, 0, NULL, NULL)))
     { goto exit; }
diff --git a/tests/test_lobject.py b/tests/test_lobject.py
index b3e3880..ef27a88 100755
--- a/tests/test_lobject.py
+++ b/tests/test_lobject.py
@@ -25,7 +25,7 @@
 import os
 import shutil
 import tempfile
-from testutils import unittest, decorate_all_tests
+from testutils import unittest, decorate_all_tests, skip_if_tpc_disabled
 
 import psycopg2
 import psycopg2.extensions
@@ -53,7 +53,7 @@ def skip_if_green(f):
 class LargeObjectMixin(object):
     # doesn't derive from TestCase to avoid repeating tests twice.
     def setUp(self):
-        self.conn = psycopg2.connect(tests.dsn)
+        self.conn = self.connect()
         self.lo_oid = None
         self.tmpdir = None
 
@@ -74,6 +74,9 @@ class LargeObjectMixin(object):
                 lo.unlink()
         self.conn.close()
 
+    def connect(self):
+        return psycopg2.connect(tests.dsn)
+
 
 class LargeObjectTests(LargeObjectMixin, unittest.TestCase):
     def test_create(self):
@@ -312,6 +315,28 @@ class LargeObjectTests(LargeObjectMixin, unittest.TestCase):
         self.assertTrue(os.path.exists(filename))
         self.assertEqual(open(filename, "rb").read(), "some data")
 
+    @skip_if_tpc_disabled
+    def test_read_after_tpc_commit(self):
+        self.conn.tpc_begin('test_lobject')
+        lo = self.conn.lobject()
+        self.lo_oid = lo.oid
+        self.conn.tpc_commit()
+
+        self.assertRaises(psycopg2.ProgrammingError, lo.read, 5)
+
+    @skip_if_tpc_disabled
+    def test_read_after_tpc_prepare(self):
+        self.conn.tpc_begin('test_lobject')
+        lo = self.conn.lobject()
+        self.lo_oid = lo.oid
+        self.conn.tpc_prepare()
+
+        try:
+            self.assertRaises(psycopg2.ProgrammingError, lo.read, 5)
+        finally:
+            self.conn.tpc_commit()
+
+
 decorate_all_tests(LargeObjectTests, skip_if_no_lo)
 decorate_all_tests(LargeObjectTests, skip_if_green)