Fix support for two phase commit in pg8000 dialect

The postgresql base dialect has problems with two-phase commit because
there isn't a standard way of handling autocommit in DBAPI. This commit
modifies the pg8000 dialect to use the DBAPI tpc extension, which is
supported by pg8000 as of version 1.9.11.

1 files changed, 19 insertions, 0 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/pg8000.py b/lib/sqlalchemy/dialects/postgresql/pg8000.py
index dc5ed6e73..dd63d7ea8 100644
--- a/lib/sqlalchemy/dialects/postgresql/pg8000.py
+++ b/lib/sqlalchemy/dialects/postgresql/pg8000.py
@@ -167,4 +167,23 @@ class PGDialect_pg8000(PGDialect):
                 (level, self.name, ", ".join(self._isolation_lookup))
                 )
 
+    def do_begin_twophase(self, connection, xid):
+        print("begin twophase", xid)
+        connection.connection.tpc_begin((0, xid, ''))
+
+    def do_prepare_twophase(self, connection, xid):
+        print("prepare twophase", xid)
+        connection.connection.tpc_prepare()
+
+    def do_rollback_twophase(
+            self, connection, xid, is_prepared=True, recover=False):
+        connection.connection.tpc_rollback((0, xid, ''))
+
+    def do_commit_twophase(
+            self, connection, xid, is_prepared=True, recover=False):
+        connection.connection.tpc_commit((0, xid, ''))
+
+    def do_recover_twophase(self, connection):
+        return [row[1] for row in connection.connection.tpc_recover()]
+
 dialect = PGDialect_pg8000