Failing test case for database.transaction() as a decorator (#192)

* Failing test case for database.transaction() as a decorator

* Rollback #158

* Tweak .force_rollback

* Revert "Tweak .force_rollback"

This reverts commit 2835db7.

* Revert "Rollback #158"

This reverts commit c18b19c.

* Fix transction decorator

* Linting

* Version 0.3.1

Author: lovelydinosaur

databases/__init__.py

@@ -1,4 +1,4 @@
 from databases.core import Database, DatabaseURL
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 __all__ = ["Database", "DatabaseURL"]

databases/core.py

@@ -185,7 +185,7 @@ def connection(self) -> "Connection":
             return connection
 
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
-        return self.connection().transaction(force_rollback=force_rollback)
+        return Transaction(self.connection, force_rollback=force_rollback)
 
     @contextlib.contextmanager
     def force_rollback(self) -> typing.Iterator[None]:
@@ -275,7 +275,10 @@ async def iterate(
                     yield record
 
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
-        return Transaction(self, force_rollback)
+        def connection_callable() -> Connection:
+            return self
+
+        return Transaction(connection_callable, force_rollback)
 
     @property
     def raw_connection(self) -> typing.Any:
@@ -296,10 +299,13 @@ def _build_query(
 
 
 class Transaction:
-    def __init__(self, connection: Connection, force_rollback: bool) -> None:
-        self._connection = connection
+    def __init__(
+        self,
+        connection_callable: typing.Callable[[], Connection],
+        force_rollback: bool,
+    ) -> None:
+        self._connection_callable = connection_callable
         self._force_rollback = force_rollback
-        self._transaction = connection._connection.transaction()
 
     async def __aenter__(self) -> "Transaction":
         """
@@ -341,6 +347,9 @@ async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
         return wrapper
 
     async def start(self) -> "Transaction":
+        self._connection = self._connection_callable()
+        self._transaction = self._connection._connection.transaction()
+
         async with self._connection._transaction_lock:
             is_root = not self._connection._transaction_stack
             await self._connection.__aenter__()

tests/test_databases.py

@@ -451,15 +451,16 @@ async def test_transaction_decorator(database_url):
     """
     Ensure that @database.transaction() is supported.
     """
-    async with Database(database_url, force_rollback=True) as database:
+    database = Database(database_url, force_rollback=True)
 
-        @database.transaction()
-        async def insert_data(raise_exception):
-            query = notes.insert().values(text="example", completed=True)
-            await database.execute(query)
-            if raise_exception:
-                raise RuntimeError()
+    @database.transaction()
+    async def insert_data(raise_exception):
+        query = notes.insert().values(text="example", completed=True)
+        await database.execute(query)
+        if raise_exception:
+            raise RuntimeError()
 
+    async with database:
         with pytest.raises(RuntimeError):
             await insert_data(raise_exception=True)