Add 'database.force_rollback()' context manager. (#189)

* Add 'database.force_rollback()' context manager

* Add .force_rollback test case

* Typing

Author: lovelydinosaur

databases/core.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import functools
 import logging
 import sys
@@ -186,6 +187,15 @@ def connection(self) -> "Connection":
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
         return self.connection().transaction(force_rollback=force_rollback)
 
+    @contextlib.contextmanager
+    def force_rollback(self) -> typing.Iterator[None]:
+        initial = self._force_rollback
+        self._force_rollback = True
+        try:
+            yield
+        finally:
+            self._force_rollback = initial
+
 
 class Connection:
     def __init__(self, backend: DatabaseBackend) -> None:

tests/test_databases.py

@@ -337,6 +337,28 @@ async def test_rollback_isolation(database_url):
         assert len(results) == 0
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_rollback_isolation_with_contextmanager(database_url):
+    """
+    Ensure that `database.force_rollback()` provides strict isolation.
+    """
+
+    database = Database(database_url)
+
+    with database.force_rollback():
+        async with database:
+            # Perform some INSERT operations on the database.
+            query = notes.insert().values(text="example1", completed=True)
+            await database.execute(query)
+
+        async with database:
+            # Ensure INSERT operations have been rolled back.
+            query = notes.select()
+            results = await database.fetch_all(query=query)
+            assert len(results) == 0
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_transaction_commit(database_url):