PATCH: two phase commit

Index: psycopg/connection_type.c
===================================================================
--- psycopg/connection_type.c	(revision 926)
+++ psycopg/connection_type.c	(working copy)
@@ -33,6 +33,7 @@
 #include "psycopg/psycopg.h"
 #include "psycopg/connection.h"
 #include "psycopg/cursor.h"
+#include "psycopg/pqpath.h"
 
 /** DBAPI methods **/
 
@@ -147,6 +148,67 @@
 
 #ifdef PSYCOPG_EXTENSIONS
 
+#define psyco_conn_prepare_transaction_doc \
+"prepare_transaction(xid) -- Prepare a transaction for 2-phase commit."
+
+static PyObject *
+psyco_conn_prepare_transaction(connectionObject *self, PyObject *args)
+{
+    const char *xid;
+
+    EXC_IF_CONN_CLOSED(self);
+
+    if (!PyArg_ParseTuple(args, "s", &xid))
+        return NULL;
+
+    self->mark++;
+    if (pq_prepare_transaction(self, xid) < 0)
+        return NULL;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+#define psyco_conn_commit_prepared_doc \
+"commit_prepared(xid) -- Commit a previously prepared transaction."
+
+static PyObject *
+psyco_conn_commit_prepared(connectionObject *self, PyObject *args)
+{
+    const char *xid;
+
+    EXC_IF_CONN_CLOSED(self);
+
+    if (!PyArg_ParseTuple(args, "s", &xid))
+        return NULL;
+
+    if (pq_commit_prepared(self, xid) < 0)
+        return NULL;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+#define psyco_conn_rollback_prepared_doc \
+"rollback_prepared(xid) -- Roll back a previously prepared transaction."
+
+static PyObject *
+psyco_conn_rollback_prepared(connectionObject *self, PyObject *args)
+{
+    const char *xid;
+
+    EXC_IF_CONN_CLOSED(self);
+
+    if (!PyArg_ParseTuple(args, "s", &xid))
+        return NULL;
+
+    if (pq_rollback_prepared(self, xid) < 0)
+        return NULL;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 /* set_isolation_level method - switch connection isolation level */
 
 #define psyco_conn_set_isolation_level_doc \
@@ -168,8 +230,6 @@
     }
 
     if (conn_switch_isolation_level(self, level) < 0) {
-        PyErr_SetString(OperationalError,
-                        PQerrorMessage(self->pgconn));
         return NULL;
     }
 
@@ -255,6 +315,12 @@
     {"rollback", (PyCFunction)psyco_conn_rollback,
      METH_VARARGS, psyco_conn_rollback_doc},
 #ifdef PSYCOPG_EXTENSIONS
+    {"prepare_transaction", (PyCFunction)psyco_conn_prepare_transaction,
+     METH_VARARGS, psyco_conn_prepare_transaction_doc},
+    {"commit_prepared", (PyCFunction)psyco_conn_commit_prepared,
+     METH_VARARGS, psyco_conn_commit_prepared_doc},
+    {"rollback_prepared", (PyCFunction)psyco_conn_rollback_prepared,
+     METH_VARARGS, psyco_conn_rollback_prepared_doc},
     {"set_isolation_level", (PyCFunction)psyco_conn_set_isolation_level,
      METH_VARARGS, psyco_conn_set_isolation_level_doc},
     {"set_client_encoding", (PyCFunction)psyco_conn_set_client_encoding,
Index: psycopg/pqpath.c
===================================================================
--- psycopg/pqpath.c	(revision 926)
+++ psycopg/pqpath.c	(working copy)
@@ -506,6 +506,132 @@
     return retvalue;
 }
 
+/* run_prepared_txn_command - a prepared transaction helper. */
+static int
+run_prepared_txn_command(connectionObject *conn, const char *command,
+                         const char *xid)
+{
+    int command_len, xid_len, buffer_len, err = 0, retvalue = 0;
+    PGresult *pgres = NULL;
+    char *error = NULL, *buffer = NULL, *tmp;
+
+    Dprintf("run_prepared_txn_command: pgconn = %p, command = %s, xid = %s",
+            conn->pgconn, command, xid);
+    command_len = strlen(command);
+    xid_len = strlen(xid);
+
+    /* PQescapeStringConn() requires that we allocate at least twice
+       the length of the escaped string plus one. */
+    buffer_len = command_len + xid_len * 2 + 5;
+    buffer = PyMem_Malloc(buffer_len);
+    if (!buffer) {
+        PyErr_NoMemory();
+        retvalue = -1;
+        goto end;
+    }
+
+    tmp = buffer;
+    strcpy(tmp, command);
+    tmp += command_len;
+    strcpy(tmp, " '");
+    tmp += 2;
+
+#if PG_MAJOR_VERSION > 8 || \
+ (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION > 1) || \
+ (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION == 1 && PG_PATCH_VERSION >= 4)
+    tmp += PQescapeStringConn(conn->pgconn, tmp, xid, xid_len, &err);
+#else
+    tmp += PQescapeString(tmp, xid, xid_len);
+#endif
+
+    if (err != 0) {
+        pq_raise(conn, NULL, NULL, NULL, NULL);
+        retvalue = -1;
+        goto end;
+    }
+    strcpy(tmp, "'");
+
+    Dprintf("run_prepared_txn_command: query is \"%s\"", buffer);
+
+    Py_BEGIN_ALLOW_THREADS;
+    pthread_mutex_lock(&conn->lock);
+
+    pq_clear_async(conn);
+    retvalue = pq_execute_command_locked(conn, buffer, &pgres, &error);
+    /* PREPARE TRANSACTION closes off the current transaction whether
+     * it passes or fails.  COMMIT PREPARED and ROLLBACK PREPARED are
+     * both run outside of a transaction.
+     */
+    conn->status = CONN_STATUS_READY;
+
+    pthread_mutex_unlock(&conn->lock);
+    Py_END_ALLOW_THREADS;
+
+    if (retvalue < 0)
+        pq_complete_error(conn, &pgres, &error);
+
+ end:
+    if (buffer)
+        PyMem_Free(buffer);
+    return retvalue;
+}
+
+/* pq_prepare_transaction - prepare a two-phase transaction */
+int
+pq_prepare_transaction(connectionObject *conn, const char *xid)
+{
+    Dprintf("pq_prepare_transaction: pgconn = %p, isolevel = %ld, "
+            "status = %d, xid = %s",
+            conn->pgconn, conn->isolation_level, conn->status, xid);
+
+    if (conn->isolation_level == 0 || conn->status != CONN_STATUS_BEGIN) {
+        Dprintf("pq_prepare_transaction: no transaction to prepare");
+        psyco_set_error(ProgrammingError, NULL, "No transaction to prepare",
+                        NULL, NULL);
+        return -1;
+    }
+
+    return run_prepared_txn_command(conn, "PREPARE TRANSACTION", xid);
+}
+
+/* pq_commit_prepared - commit a two-phase transaction */
+int
+pq_commit_prepared(connectionObject *conn, const char *xid)
+{
+    Dprintf("pq_commit_prepared: pgconn = %p, isolevel = %ld, "
+            "status = %d, xid = %s",
+            conn->pgconn, conn->isolation_level, conn->status, xid);
+
+    if (conn->isolation_level != 0 && conn->status != CONN_STATUS_READY) {
+        Dprintf("pq_commit_prepared: inside a transaction");
+        psyco_set_error(ProgrammingError, NULL,
+                        "commit_prepared cannot run inside a transaction",
+                        NULL, NULL);
+        return -1;
+    }
+
+    return run_prepared_txn_command(conn, "COMMIT PREPARED", xid);
+}
+
+/* pq_rollback_prepared - roll back a two-phase transaction */
+int
+pq_rollback_prepared(connectionObject *conn, const char *xid)
+{
+    Dprintf("pq_rollback_prepared: pgconn = %p, isolevel = %ld, "
+            "status = %d, xid = %s",
+            conn->pgconn, conn->isolation_level, conn->status, xid);
+
+    if (conn->isolation_level != 0 && conn->status != CONN_STATUS_READY) {
+        Dprintf("pq_rollback_prepared: inside a transaction");
+        psyco_set_error(ProgrammingError, NULL,
+                        "rollback_prepared cannot run inside a transaction",
+                        NULL, NULL);
+        return -1;
+    }
+
+    return run_prepared_txn_command(conn, "ROLLBACK PREPARED", xid);
+}
+
 /* pq_is_busy - consume input and return connection status
 
    a status of 1 means that a call to pq_fetch will block, while a status of 0
Index: psycopg/pqpath.h
===================================================================
--- psycopg/pqpath.h	(revision 926)
+++ psycopg/pqpath.h	(working copy)
@@ -38,6 +38,10 @@
 extern int pq_abort(connectionObject *conn);
 extern int pq_is_busy(connectionObject *conn);
 
+extern int pq_prepare_transaction(connectionObject *conn, const char *xid);
+extern int pq_commit_prepared(connectionObject *conn, const char *xid);
+extern int pq_rollback_prepared(connectionObject *conn, const char *xid);
+
 extern void pq_set_critical(connectionObject *conn, const char *msg);
 extern PyObject *pq_resolve_critical(connectionObject *conn, int close);
 
Index: ChangeLog
===================================================================
--- ChangeLog	(revision 926)
+++ ChangeLog	(working copy)
@@ -1,5 +1,20 @@
 2008-01-17  James Henstridge  <[email protected]>
 
+	* tests/test_transaction.py (TwoPhaseCommitTests): add some tests
+	for two phase commit.
+
+	* psycopg/connection_type.c: add prepare_transaction(),
+	commit_prepared() and rollback_prepared() methods.
+
+	* psycopg/pqpath.c (run_prepared_txn_command): a helper function
+	for prepared transaction commands.
+	(pq_prepare_transaction): new function for preparing a
+	transaction.
+	(pq_commit_prepared): new function for committing a prepared
+	transaction.
+	(pq_rollback_prepared): new function for rolling back a prepared
+	transaction.
+
 	* ZPsycopgDA/DA.py (Connection.__init__): Default the encoding to
 	UTF-8, fixing bug #190.
 	(App.ImageFile): simplify ImageFile import using patch from
Index: tests/test_transaction.py
===================================================================
--- tests/test_transaction.py	(revision 926)
+++ tests/test_transaction.py	(working copy)
@@ -72,6 +72,119 @@
         self.assertEqual(curs.fetchone()[0], 1)
 
 
+class TwoPhaseCommitTests(unittest.TestCase):
+    """Tests of two phase commit functionality."""
+
+    def setUp(self):
+        self.conn = psycopg2.connect("dbname=%s" % tests.dbname)
+        self.conn.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
+
+        # Create table for test:
+        curs = self.conn.cursor()
+        try:
+            curs.execute("DROP TABLE table1")
+            self.conn.commit()
+        except psycopg2.DatabaseError:
+            self.conn.rollback()
+        curs.execute('''
+            CREATE TABLE table1 (
+              id int PRIMARY KEY
+            )''')
+
+        # pick two unused transaction IDs for the tests:
+        existing_xids = self.get_xids()
+        self.test_xid = None
+        i = 0
+        while self.test_xid is None:
+            xid = 'test_xid_%d' % i
+            if xid not in existing_xids:
+                self.test_xid = xid
+            i += 1
+        self.conn.commit()
+
+    def tearDown(self):
+        self.conn.rollback()
+        try:
+            self.conn.rollback_prepared(self.test_xid)
+        except psycopg2.DatabaseError:
+            pass
+        curs = self.conn.cursor()
+        curs.execute('DROP TABLE table1')
+        self.conn.commit()
+        self.conn.close()
+
+    def get_xids(self):
+        curs = self.conn.cursor()
+        curs.execute('SELECT gid FROM pg_prepared_xacts')
+        return [row[0] for row in curs.fetchall()]
+
+    def test_commit_prepared(self):
+        # Prepare a transaction:
+        curs = self.conn.cursor()
+        curs.execute('INSERT INTO table1 VALUES (42)')
+        self.assertEqual(self.conn.status, STATUS_BEGIN)
+        self.conn.prepare_transaction(self.test_xid)
+        self.assertEqual(self.conn.status, STATUS_READY)
+
+        # The changes from the transaction are not visible:
+        curs.execute('SELECT COUNT(*) FROM table1 WHERE id = 42')
+        self.assertEqual(curs.fetchone()[0], 0)
+        self.conn.rollback()
+
+        # After committing the prepared transaction, the row is visible:
+        self.conn.commit_prepared(self.test_xid)
+        self.assertEqual(self.conn.status, STATUS_READY)
+        curs.execute('SELECT COUNT(*) FROM table1 WHERE id = 42')
+        self.assertEqual(curs.fetchone()[0], 1)
+
+    def test_rollback_prepared(self):
+        # Prepare a transaction:
+        curs = self.conn.cursor()
+        curs.execute('INSERT INTO table1 VALUES (42)')
+        self.assertEqual(self.conn.status, STATUS_BEGIN)
+        self.conn.prepare_transaction(self.test_xid)
+        self.assertEqual(self.conn.status, STATUS_READY)
+
+        # The changes from the transaction are not visible:
+        curs.execute('SELECT COUNT(*) FROM table1 WHERE id = 42')
+        self.assertEqual(curs.fetchone()[0], 0)
+        self.conn.rollback()
+
+        # Now roll back the transaction.  The changes are lost.
+        self.conn.rollback_prepared(self.test_xid)
+        self.assertEqual(self.conn.status, STATUS_READY)
+        curs.execute('SELECT COUNT(*) FROM table1 WHERE id = 42')
+        self.assertEqual(curs.fetchone()[0], 0)
+
+    def test_commit_unknown(self):
+        self.assertRaises(psycopg2.ProgrammingError,
+                          self.conn.commit_prepared, self.test_xid)
+
+    def test_rollback_unknown(self):
+        self.assertRaises(psycopg2.ProgrammingError,
+                          self.conn.rollback_prepared, self.test_xid)
+
+    def test_duplicate_transaction_id(self):
+        curs = self.conn.cursor()
+        curs.execute('INSERT INTO table1 VALUES (42)')
+        self.assertEqual(self.conn.status, STATUS_BEGIN)
+        self.conn.prepare_transaction(self.test_xid)
+        self.assertEqual(self.conn.status, STATUS_READY)
+
+        # Now try to reuse the transaction ID, which results in an
+        # error:
+        curs.execute('INSERT INTO table1 VALUES (43)')
+        self.assertEqual(self.conn.status, STATUS_BEGIN)
+        self.assertRaises(psycopg2.ProgrammingError,
+                          self.conn.prepare_transaction, self.test_xid)
+
+        # The changes from the failed prepared transaction have been
+        # discarded:
+        self.assertEqual(self.conn.status, STATUS_READY)
+        curs.execute('SELECT COUNT(*) FROM table1 WHERE id = 43')
+        self.assertEqual(curs.fetchone()[0], 0)
+
+
 class DeadlockSerializationTestCase(unittest.TestCase):
     """Test deadlock and serialization failure errors."""