From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 24 Nov 2020 12:26:18 +0000 (+0000)
Subject: Make sure transaction blocks cannot be used more than once
X-Git-Tag: 3.0.dev0~311
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f33ed76766984b6725af7f3640f52fd5c398154e;p=thirdparty%2Fpsycopg.git

Make sure transaction blocks cannot be used more than once

See #12
---

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()