From 6b08ae70253d2fbd867aa6fca794e5648aa83a0e Mon Sep 17 00:00:00 2001
From: Federico Caselli <cfederico87@gmail.com>
Date: Wed, 27 Oct 2021 22:52:39 +0200
Subject: [PATCH] Update psycopg2 dialect to use the DBAPI interface to execute
 two phase transactions

Fixes: #7238
Change-Id: Ie4f5cf59d29b5bfc142ec2dfdecffb896ee7d708
---
 doc/build/changelog/unreleased_20/7238.rst    |  7 +++
 .../dialects/postgresql/psycopg2.py           | 46 +++++++++++++++----
 .../dialects/postgresql/psycopg2cffi.py       |  5 +-
 3 files changed, 48 insertions(+), 10 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_20/7238.rst

diff --git a/doc/build/changelog/unreleased_20/7238.rst b/doc/build/changelog/unreleased_20/7238.rst
new file mode 100644
index 0000000000..87ef43c2b6
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/7238.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: postgresql, psycopg2
+    :tickets: 7238
+
+    Update psycopg2 dialect to use the DBAPI interface to execute
+    two phase transactions. Previously SQL commands were execute
+    to handle this kind of transactions.
diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 4143dd041d..7512ab9b52 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -783,13 +783,13 @@ class PGDialect_psycopg2(PGDialect):
 
         return psycopg2
 
-    @classmethod
+    @util.memoized_property
     def _psycopg2_extensions(cls):
         from psycopg2 import extensions
 
         return extensions
 
-    @classmethod
+    @util.memoized_property
     def _psycopg2_extras(cls):
         from psycopg2 import extras
 
@@ -797,7 +797,7 @@ class PGDialect_psycopg2(PGDialect):
 
     @util.memoized_property
     def _isolation_lookup(self):
-        extensions = self._psycopg2_extensions()
+        extensions = self._psycopg2_extensions
         return {
             "AUTOCOMMIT": extensions.ISOLATION_LEVEL_AUTOCOMMIT,
             "READ COMMITTED": extensions.ISOLATION_LEVEL_READ_COMMITTED,
@@ -853,8 +853,8 @@ class PGDialect_psycopg2(PGDialect):
             return True
 
     def on_connect(self):
-        extras = self._psycopg2_extras()
-        extensions = self._psycopg2_extensions()
+        extras = self._psycopg2_extras
+        extensions = self._psycopg2_extensions
 
         fns = []
         if self.client_encoding is not None:
@@ -947,7 +947,7 @@ class PGDialect_psycopg2(PGDialect):
                 kwargs = {"page_size": self.executemany_values_page_size}
             else:
                 kwargs = {}
-            xtras = self._psycopg2_extras()
+            xtras = self._psycopg2_extras
             context._psycopg2_fetched_rows = xtras.execute_values(
                 cursor,
                 statement,
@@ -962,15 +962,45 @@ class PGDialect_psycopg2(PGDialect):
                 kwargs = {"page_size": self.executemany_batch_page_size}
             else:
                 kwargs = {}
-            self._psycopg2_extras().execute_batch(
+            self._psycopg2_extras.execute_batch(
                 cursor, statement, parameters, **kwargs
             )
         else:
             cursor.executemany(statement, parameters)
 
+    def do_begin_twophase(self, connection, xid):
+        connection.connection.tpc_begin(xid)
+
+    def do_prepare_twophase(self, connection, xid):
+        connection.connection.tpc_prepare()
+
+    def _do_twophase(self, dbapi_conn, operation, xid, recover=False):
+        if recover:
+            if dbapi_conn.status != self._psycopg2_extensions.STATUS_READY:
+                dbapi_conn.rollback()
+            operation(xid)
+        else:
+            operation()
+
+    def do_rollback_twophase(
+        self, connection, xid, is_prepared=True, recover=False
+    ):
+        dbapi_conn = connection.connection.dbapi_connection
+        self._do_twophase(
+            dbapi_conn, dbapi_conn.tpc_rollback, xid, recover=recover
+        )
+
+    def do_commit_twophase(
+        self, connection, xid, is_prepared=True, recover=False
+    ):
+        dbapi_conn = connection.connection.dbapi_connection
+        self._do_twophase(
+            dbapi_conn, dbapi_conn.tpc_commit, xid, recover=recover
+        )
+
     @util.memoized_instancemethod
     def _hstore_oids(self, conn):
-        extras = self._psycopg2_extras()
+        extras = self._psycopg2_extras
         if hasattr(conn, "dbapi_connection"):
             conn = conn.dbapi_connection
         oids = extras.HstoreAdapter.get_oids(conn)
diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2cffi.py b/lib/sqlalchemy/dialects/postgresql/psycopg2cffi.py
index 5be52a8707..05f1837c60 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2cffi.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2cffi.py
@@ -23,6 +23,7 @@ is as per ``psycopg2``.
 
 """  # noqa
 from .psycopg2 import PGDialect_psycopg2
+from ... import util
 
 
 class PGDialect_psycopg2cffi(PGDialect_psycopg2):
@@ -46,12 +47,12 @@ class PGDialect_psycopg2cffi(PGDialect_psycopg2):
     def dbapi(cls):
         return __import__("psycopg2cffi")
 
-    @classmethod
+    @util.memoized_property
     def _psycopg2_extensions(cls):
         root = __import__("psycopg2cffi", fromlist=["extensions"])
         return root.extensions
 
-    @classmethod
+    @util.memoized_property
     def _psycopg2_extras(cls):
         root = __import__("psycopg2cffi", fromlist=["extras"])
         return root.extras
-- 
2.47.2