fix concurrency issues of parallel transactions (#327)

databases/core.py

@@ -5,7 +5,7 @@
 import sys
 import typing
 from types import TracebackType
-from urllib.parse import SplitResult, parse_qsl, urlsplit, unquote
+from urllib.parse import SplitResult, parse_qsl, unquote, urlsplit
 
 from sqlalchemy import text
 from sqlalchemy.sql import ClauseElement
@@ -14,8 +14,10 @@
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
 
 if sys.version_info >= (3, 7):  # pragma: no cover
+    import contextvars as contextvars
     from contextvars import ContextVar
 else:  # pragma: no cover
+    import aiocontextvars as contextvars
     from aiocontextvars import ContextVar
 
 try:  # pragma: no cover
@@ -173,21 +175,35 @@ async def iterate(
             async for record in connection.iterate(query, values):
                 yield record
 
+    def _new_connection(self) -> "Connection":
+        connection = Connection(self._backend)
+        self._connection_context.set(connection)
+        return connection
+
     def connection(self) -> "Connection":
         if self._global_connection is not None:
             return self._global_connection
 
         try:
             return self._connection_context.get()
         except LookupError:
-            connection = Connection(self._backend)
-            self._connection_context.set(connection)
-            return connection
+            return self._new_connection()
 
     def transaction(
         self, *, force_rollback: bool = False, **kwargs: typing.Any
     ) -> "Transaction":
-        return Transaction(self.connection, force_rollback=force_rollback, **kwargs)
+        try:
+            connection = self._connection_context.get()
+            is_root = not connection._transaction_stack
+            if is_root:
+                newcontext = contextvars.copy_context()
+                get_conn = lambda: newcontext.run(self._new_connection)
+            else:
+                get_conn = self.connection
+        except LookupError:
+            get_conn = self.connection
+
+        return Transaction(get_conn, force_rollback=force_rollback, **kwargs)
 
     @contextlib.contextmanager
     def force_rollback(self) -> typing.Iterator[None]:

tests/test_database_url.py

@@ -1,7 +1,9 @@
-from databases import DatabaseURL
 from urllib.parse import quote
+
 import pytest
 
+from databases import DatabaseURL
+
 
 def test_database_url_repr():
     u = DatabaseURL("postgresql://localhost/name")

tests/test_databases.py

@@ -800,7 +800,7 @@ async def test_queries_with_expose_backend_connection(database_url):
     """
     async with Database(database_url) as database:
         async with database.connection() as connection:
-            async with database.transaction(force_rollback=True):
+            async with connection.transaction(force_rollback=True):
                 # Get the raw connection
                 raw_connection = connection.raw_connection
 
@@ -996,3 +996,21 @@ async def test_column_names(database_url, select_query):
             assert sorted(results[0].keys()) == ["completed", "id", "text"]
             assert results[0]["text"] == "example1"
             assert results[0]["completed"] == True
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_parallel_transactions(database_url):
+    """
+    Test parallel transaction execution.
+    """
+
+    async def test_task(db):
+        async with db.transaction():
+            await db.fetch_one("SELECT 1")
+
+    async with Database(database_url) as database:
+        await database.fetch_one("SELECT 1")
+
+        tasks = [test_task(database) for i in range(4)]
+        await asyncio.gather(*tasks)