Update psycopg2 dialect to use the DBAPI interface to execute

two phase transactions

Fixes: #7238
Change-Id: Ie4f5cf59d29b5bfc142ec2dfdecffb896ee7d708

diff --git a/doc/build/changelog/unreleased_20/7238.rst b/doc/build/changelog/unreleased_20/7238.rst
new file mode 100644 (file)
index 0000000..87ef43c
--- /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 4143dd041d65cc88f1bf14812ad900f27303d147..7512ab9b524061910ce91800c9743006fa4d0162 100644 (file)
--- 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 5be52a8707b2edff49cc3ee9a9fb33bc93f57eef..05f1837c60af7c49c7f40dcafd89e280931efa75 100644 (file)
--- 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