From f33ed76766984b6725af7f3640f52fd5c398154e Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 24 Nov 2020 12:26:18 +0000
Subject: [PATCH] Make sure transaction blocks cannot be used more than once

See #12
---
 psycopg3/psycopg3/transaction.py | 6 ++++--
 tests/test_transaction.py        | 9 +++++++++
 tests/test_transaction_async.py  | 9 +++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/psycopg3/psycopg3/transaction.py b/psycopg3/psycopg3/transaction.py
index 8130e44ac..14c6adaba 100644
--- a/psycopg3/psycopg3/transaction.py
+++ b/psycopg3/psycopg3/transaction.py
@@ -75,8 +75,10 @@ class BaseTransaction(Generic[ConnectionType]):
         return f"{self.__class__.__qualname__}({', '.join(args)})"
 
     def _enter_commands(self) -> List[str]:
-        assert self._yolo
-        self._yolo = False
+        if not self._yolo:
+            raise TypeError("transaction blocks cannot be use more than once")
+        else:
+            self._yolo = False
 
         self._outer_transaction = (
             self._conn.pgconn.transaction_status == TransactionStatus.IDLE
diff --git a/tests/test_transaction.py b/tests/test_transaction.py
index 86652b954..1caa2c289 100644
--- a/tests/test_transaction.py
+++ b/tests/test_transaction.py
@@ -106,6 +106,15 @@ def test_exposes_savepoint_name(conn):
             tx.savepoint_name = "bar"
 
 
+def test_cant_reenter(conn):
+    with conn.transaction() as tx:
+        pass
+
+    with pytest.raises(TypeError):
+        with tx:
+            pass
+
+
 def test_begins_on_enter(conn):
     """Transaction does not begin until __enter__() is called."""
     tx = conn.transaction()
diff --git a/tests/test_transaction_async.py b/tests/test_transaction_async.py
index 8cae48e70..f9cbd420d 100644
--- a/tests/test_transaction_async.py
+++ b/tests/test_transaction_async.py
@@ -50,6 +50,15 @@ async def test_exposes_savepoint_name(aconn):
             tx.savepoint_name = "bar"
 
 
+async def test_cant_reenter(aconn):
+    async with aconn.transaction() as tx:
+        pass
+
+    with pytest.raises(TypeError):
+        async with tx:
+            pass
+
+
 async def test_begins_on_enter(aconn):
     """Transaction does not begin until __enter__() is called."""
     tx = aconn.transaction()
-- 
2.47.2