Handle correct disconnect for failed PostgreSQL transactions

Author: ansipunk

README.md

@@ -105,7 +105,7 @@ Pass database URLs for those you want to run the tests against. Comma separated
 list.
 
 ```bash
-BASED_TEST_DB_URLS='postgresql://postgres:postgres@localhost:5432/postgres,mysql://root:[email protected]:3306/mariadb' make test`
+BASED_TEST_DB_URLS='postgresql://postgres:postgres@localhost:5432/postgres,mysql://root:[email protected]:3306/mariadb' make test
 ```
 
 ## TODO

based/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.1"
+__version__ = "0.7.0"
 
 from based.backends import Session
 from based.database import Database

based/backends/__init__.py

@@ -118,6 +118,7 @@ class Session:
     _conn: typing.Any
     _dialect: Dialect
     _transaction_stack: typing.List[str]
+    transaction_failed: bool
 
     def __init__(  # noqa: D107
         self,
@@ -127,6 +128,7 @@ def __init__(  # noqa: D107
         self._conn = conn
         self._dialect = dialect
         self._transaction_stack = []
+        self.transaction_failed = False
 
     async def _execute(
         self,
@@ -154,7 +156,11 @@ async def _execute(
             cursor:
                 A cursor returned by the database driver after executing the query.
         """
-        return await self._conn.execute(query, params)
+        try:
+            return await self._conn.execute(query, params)
+        except Exception:
+            self.transaction_failed = True
+            raise
 
     def _compile_query(
         self,

based/backends/postgresql.py

@@ -84,14 +84,20 @@ async def _session(self) -> typing.AsyncGenerator["Session", None]:
                 await session.cancel_transaction()
                 raise
             else:
-                await session.commit_transaction()
+                if session.transaction_failed:
+                    await session.cancel_transaction()
+                else:
+                    await session.commit_transaction()
         else:
             try:
                 yield session
             except Exception:
                 await connection.rollback()
                 raise
             else:
-                await connection.commit()
+                if session.transaction_failed:
+                    await connection.rollback()
+                else:
+                    await connection.commit()
             finally:
                 await self._pool.putconn(connection)

tests/test_backend.py

@@ -1,3 +1,4 @@
+import contextlib
 import typing
 
 import pytest
@@ -33,7 +34,9 @@ async def test_database_force_rollback_with_lock(
     title, year = gen_movie()
 
     async with based.Database(
-        database_url, force_rollback=True, use_lock=True,
+        database_url,
+        force_rollback=True,
+        use_lock=True,
     ) as database:
         async with database.session() as session:
             query = table.insert().values(title=title, year=year)
@@ -85,3 +88,17 @@ async def test_abstract_backend(database_url: str):
 
     with pytest.raises(NotImplementedError):
         await backend.disconnect()
+
+
+async def test_disconnect_with_failed_transaction_force_rollback(database_url: str):
+    async with based.Database(database_url, force_rollback=True) as database:
+        async with database.session() as session:
+            with contextlib.suppress(Exception):
+                await session.execute("SELECT 1 FROM nonexistent;")
+
+
+async def test_disconnect_with_failed_transaction_no_force_rollback(database_url: str):
+    async with based.Database(database_url, force_rollback=False) as database:
+        async with database.session() as session:
+            with contextlib.suppress(Exception):
+                await session.execute("SELECT 1 FROM nonexistent;")