Don't run ROLLBACK when the connection is closed. (#778)

* Don't run ROLLBACK when the connection is closed.

This can be caused when a query times out while running, for example,
and the connection is closed as a result (as opposed to cancelling the
query, since PR #570). In this case, we would rather not emit the
ROLLBACK (the connection is already closed, so the transaction is over
anyway), rather than raising an exception when trying to use a
connection which is already closed.

See issue #777.

* add tests for cancelled queries in transaction

* Update test_transaction.py

* Update test_sa_transaction.py

* Update connection.py

Co-authored-by: Yury Pliner <[email protected]>

Author: brianmaissy

aiopg/sa/connection.py

@@ -211,6 +211,10 @@ async def _commit_impl(self):
             self._transaction = None
 
     async def _rollback_impl(self):
+        if self._connection.closed:
+            self._transaction = None
+            return
+
         cur = await self._get_cursor()
         try:
             await cur.execute('ROLLBACK')
@@ -253,6 +257,10 @@ async def _savepoint_impl(self, name=None):
             cur.close()
 
     async def _rollback_to_savepoint_impl(self, name, parent):
+        if self._connection.closed:
+            self._transaction = parent
+            return
+
         cur = await self._get_cursor()
         try:
             await cur.execute(f'ROLLBACK TO SAVEPOINT {name}')

aiopg/transaction.py

@@ -127,13 +127,15 @@ async def commit(self):
 
     async def rollback(self):
         self._check_commit_rollback()
-        await self._cur.execute(self._isolation.rollback())
+        if not self._cur.closed:
+            await self._cur.execute(self._isolation.rollback())
         self._is_begin = False
 
     async def rollback_savepoint(self):
         self._check_release_rollback()
-        await self._cur.execute(
-            self._isolation.rollback_savepoint(self._unique_id))
+        if not self._cur.closed:
+            await self._cur.execute(
+                self._isolation.rollback_savepoint(self._unique_id))
         self._unique_id = None
 
     async def release_savepoint(self):

tests/test_sa_transaction.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest import mock
 
 import pytest
@@ -411,3 +412,67 @@ async def test_transaction_mode(connect):
     res1 = await conn.scalar(select([func.count()]).select_from(tbl))
     assert 5 == res1
     await tr8.commit()
+
+
+async def test_timeout_in_transaction_context_manager(make_engine):
+    engine = await make_engine(timeout=1)
+    with pytest.raises(asyncio.TimeoutError):
+        async with engine.acquire() as connection:
+            async with connection.begin():
+                await connection.execute("SELECT pg_sleep(10)")
+
+    engine.terminate()
+    await engine.wait_closed()
+
+
+async def test_timeout_in_nested_transaction_context_manager(make_engine):
+    engine = await make_engine(timeout=1)
+    with pytest.raises(asyncio.TimeoutError):
+        async with engine.acquire() as connection:
+            async with connection.begin():
+                async with connection.begin_nested():
+                    await connection.execute("SELECT pg_sleep(10)")
+
+    engine.terminate()
+    await engine.wait_closed()
+
+
+async def test_cancel_in_transaction_context_manager(make_engine, loop):
+    engine = await make_engine()
+
+    with pytest.raises(asyncio.CancelledError):
+        async with engine.acquire() as connection:
+            async with connection.begin():
+                task = loop.create_task(
+                    connection.execute("SELECT pg_sleep(10)"))
+
+                async def cancel_soon():
+                    await asyncio.sleep(1)
+                    task.cancel()
+
+                loop.create_task(cancel_soon())
+                await task
+
+    engine.terminate()
+    await engine.wait_closed()
+
+
+async def test_cancel_in_savepoint_context_manager(make_engine, loop):
+    engine = await make_engine()
+
+    with pytest.raises(asyncio.CancelledError):
+        async with engine.acquire() as connection:
+            async with connection.begin():
+                async with connection.begin_nested():
+                    task = loop.create_task(
+                        connection.execute("SELECT pg_sleep(10)"))
+
+                    async def cancel_soon():
+                        await asyncio.sleep(1)
+                        task.cancel()
+
+                    loop.create_task(cancel_soon())
+                    await task
+
+    engine.terminate()
+    await engine.wait_closed()

tests/test_transaction.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import psycopg2
 import pytest
 
@@ -181,3 +183,61 @@ async def test_transaction_point_oldstyle(engine):
                        (3, 'data')]
 
         await tr.commit()
+
+
+async def test_timeout_in_transaction_context_manager(make_engine):
+    engine = await make_engine(timeout=1)
+    with pytest.raises(asyncio.TimeoutError):
+        async with engine.acquire() as connection:
+            async with Transaction(connection, IsolationLevel.read_committed):
+                await connection.execute("SELECT pg_sleep(10)")
+
+    engine.terminate()
+    await engine.wait_closed()
+
+
+async def test_timeout_in_savepoint_context_manager(make_engine):
+    engine = await make_engine(timeout=1)
+    with pytest.raises(asyncio.TimeoutError):
+        async with engine.acquire() as connection:
+            async with Transaction(
+                connection, IsolationLevel.read_committed
+            ) as transaction:
+                async with transaction.point():
+                    await connection.execute("SELECT pg_sleep(10)")
+
+    engine.terminate()
+    await engine.wait_closed()
+
+
+async def test_cancel_in_transaction_context_manager(engine, loop):
+    with pytest.raises(asyncio.CancelledError):
+        async with engine.acquire() as connection:
+            async with Transaction(connection, IsolationLevel.read_committed):
+                task = loop.create_task(
+                    connection.execute("SELECT pg_sleep(10)"))
+
+                async def cancel_soon():
+                    await asyncio.sleep(1)
+                    task.cancel()
+
+                loop.create_task(cancel_soon())
+                await task
+
+
+async def test_cancel_in_savepoint_context_manager(engine, loop):
+    with pytest.raises(asyncio.CancelledError):
+        async with engine.acquire() as connection:
+            async with Transaction(
+                connection, IsolationLevel.read_committed
+            ) as transaction:
+                async with transaction.point():
+                    task = loop.create_task(
+                        connection.execute("SELECT pg_sleep(10)"))
+
+                    async def cancel_soon():
+                        await asyncio.sleep(1)
+                        task.cancel()
+
+                    loop.create_task(cancel_soon())
+                    await task