From f90757d0912009abc28769c74b699deb23925a5f Mon Sep 17 00:00:00 2001
From: Olly Cope <olly@ollycope.com>
Date: Sat, 29 Oct 2022 16:28:40 +0000
Subject: translation handling: add checks for nested transactions

---
 yoyo/backends/base.py            |  1 +
 yoyo/backends/core/postgresql.py | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/yoyo/backends/base.py b/yoyo/backends/base.py
index 1f3c167..8ca956e 100644
--- a/yoyo/backends/base.py
+++ b/yoyo/backends/base.py
@@ -278,6 +278,7 @@ class DatabaseBackend:
         """
         Begin a new transaction
         """
+        assert not self._in_transaction
         self._in_transaction = True
         self.execute("BEGIN")
 
diff --git a/yoyo/backends/core/postgresql.py b/yoyo/backends/core/postgresql.py
index 0400ce9..d879928 100644
--- a/yoyo/backends/core/postgresql.py
+++ b/yoyo/backends/core/postgresql.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from contextlib import contextmanager
 
 from yoyo.backends.base import DatabaseBackend
@@ -26,6 +27,12 @@ class PostgresqlBackend(DatabaseBackend):
         "WHERE table_schema = :schema"
     )
 
+    @property
+    def TRANSACTION_STATUS_IDLE(self):
+        from psycopg2.extensions import TRANSACTION_STATUS_IDLE
+
+        return TRANSACTION_STATUS_IDLE
+
     def connect(self, dburi):
         kwargs = {"dbname": dburi.database}
         kwargs.update(dburi.args)
@@ -40,6 +47,16 @@ class PostgresqlBackend(DatabaseBackend):
         self.schema = kwargs.pop("schema", None)
         return self.driver.connect(**kwargs)
 
+    def transaction(self, rollback_on_exit=False):
+
+        if self.connection.info.transaction_status != self.TRANSACTION_STATUS_IDLE:
+            warnings.warn(
+                "Nested transaction requested; "
+                "this will raise an exception in some "
+                "PostgreSQL-compatible databases"
+            )
+        return super().transaction(rollback_on_exit=rollback_on_exit)
+
     @contextmanager
     def disable_transactions(self):
         with super(PostgresqlBackend, self).disable_transactions():
@@ -64,3 +81,9 @@ class PostgresqlPsycopgBackend(PostgresqlBackend):
     """
 
     driver_module = "psycopg"
+
+    @property
+    def TRANSACTION_STATUS_IDLE(self):
+        from psycopg.pq import TransactionStatus
+
+        return TransactionStatus.IDLE
-- 
cgit v1.2.1