fix: incorrect concurrent usage of connection and transaction

databases/core.py

@@ -3,15 +3,14 @@
 import functools
 import logging
 import typing
-from contextvars import ContextVar
 from types import TracebackType
 from urllib.parse import SplitResult, parse_qsl, unquote, urlsplit
 
 from sqlalchemy import text
 from sqlalchemy.sql import ClauseElement
 
 from databases.importer import import_from_string
-from databases.interfaces import DatabaseBackend, Record
+from databases.interfaces import DatabaseBackend, Record, TransactionBackend
 
 try:  # pragma: no cover
     import click
@@ -63,8 +62,8 @@ def __init__(
         assert issubclass(backend_cls, DatabaseBackend)
         self._backend = backend_cls(self.url, **self.options)
 
-        # Connections are stored as task-local state.
-        self._connection_context: ContextVar = ContextVar("connection_context")
+        # Connections are stored per asyncio task
+        self._connections: typing.Dict[asyncio.Task, Connection] = {}
 
         # When `force_rollback=True` is used, we use a single global
         # connection, within a transaction that always rolls back.
@@ -113,7 +112,10 @@ async def disconnect(self) -> None:
             self._global_transaction = None
             self._global_connection = None
         else:
-            self._connection_context = ContextVar("connection_context")
+            current_task = asyncio.current_task()
+            assert current_task is not None, "No currently running task"
+            if current_task in self._connections:
+                del self._connections[current_task]
 
         await self._backend.disconnect()
         logger.info(
@@ -187,12 +189,12 @@ 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
+        current_task = asyncio.current_task()
+        assert current_task is not None, "No currently running task"
+        if current_task not in self._connections:
+            self._connections[current_task] = Connection(self._backend)
+
+        return self._connections[current_task]
 
     def transaction(
         self, *, force_rollback: bool = False, **kwargs: typing.Any
@@ -345,6 +347,11 @@ def __init__(
         self._force_rollback = force_rollback
         self._extra_options = kwargs
 
+        # Transactions are stored per asyncio task
+        self._transactions: typing.Dict[
+            typing.Optional[asyncio.Task], TransactionBackend
+        ] = {}
+
     async def __aenter__(self) -> "Transaction":
         """
         Called when entering `async with database.transaction()`
@@ -385,31 +392,41 @@ async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
         return wrapper  # type: ignore
 
     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__()
-            await self._transaction.start(
-                is_root=is_root, extra_options=self._extra_options
-            )
-            self._connection._transaction_stack.append(self)
+        connection = self._connection_callable()
+        current_task = asyncio.current_task()
+        assert current_task is not None, "No currently running task"
+        transaction = connection._connection.transaction()
+        self._transactions[current_task] = transaction
+        async with connection._transaction_lock:
+            is_root = not connection._transaction_stack
+            await connection.__aenter__()
+            await transaction.start(is_root=is_root, extra_options=self._extra_options)
+            connection._transaction_stack.append(self)
         return self
 
     async def commit(self) -> None:
-        async with self._connection._transaction_lock:
-            assert self._connection._transaction_stack[-1] is self
-            self._connection._transaction_stack.pop()
-            await self._transaction.commit()
-            await self._connection.__aexit__()
+        connection = self._connection_callable()
+        current_task = asyncio.current_task()
+        transaction = self._transactions.get(current_task, None)
+        assert transaction is not None, "Transaction not found in current task"
+        async with connection._transaction_lock:
+            assert connection._transaction_stack[-1] is self
+            connection._transaction_stack.pop()
+            await transaction.commit()
+            await connection.__aexit__()
+            del self._transactions[current_task]
 
     async def rollback(self) -> None:
-        async with self._connection._transaction_lock:
-            assert self._connection._transaction_stack[-1] is self
-            self._connection._transaction_stack.pop()
-            await self._transaction.rollback()
-            await self._connection.__aexit__()
+        connection = self._connection_callable()
+        current_task = asyncio.current_task()
+        transaction = self._transactions.get(current_task, None)
+        assert transaction is not None, "Transaction not found in current task"
+        async with connection._transaction_lock:
+            assert connection._transaction_stack[-1] is self
+            connection._transaction_stack.pop()
+            await transaction.rollback()
+            await connection.__aexit__()
+            del self._transactions[current_task]
 
 
 class _EmptyNetloc(str):

tests/test_databases.py

@@ -2,6 +2,7 @@
 import datetime
 import decimal
 import functools
+import itertools
 import os
 import re
 from unittest.mock import MagicMock, patch
@@ -789,15 +790,16 @@ async def test_connect_and_disconnect(database_url):
 
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
-async def test_connection_context(database_url):
-    """
-    Test connection contexts are task-local.
-    """
+async def test_connection_context_same_task(database_url):
     async with Database(database_url) as database:
         async with database.connection() as connection_1:
             async with database.connection() as connection_2:
                 assert connection_1 is connection_2
 
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_connection_context_multiple_tasks(database_url):
     async with Database(database_url) as database:
         connection_1 = None
         connection_2 = None
@@ -817,9 +819,8 @@ async def get_connection_2():
                 connection_2 = connection
                 await test_complete.wait()
 
-        loop = asyncio.get_event_loop()
-        task_1 = loop.create_task(get_connection_1())
-        task_2 = loop.create_task(get_connection_2())
+        task_1 = asyncio.create_task(get_connection_1())
+        task_2 = asyncio.create_task(get_connection_2())
         while connection_1 is None or connection_2 is None:
             await asyncio.sleep(0.000001)
         assert connection_1 is not connection_2
@@ -828,6 +829,20 @@ async def get_connection_2():
         await task_2
 
 
+@pytest.mark.parametrize(
+    "database_url1,database_url2",
+    (
+        pytest.param(db1, db2, id=f"{db1} | {db2}")
+        for (db1, db2) in itertools.combinations(DATABASE_URLS, 2)
+    ),
+)
+@async_adapter
+async def test_connection_context_multiple_databases(database_url1, database_url2):
+    async with Database(database_url1) as database1:
+        async with Database(database_url2) as database2:
+            assert database1.connection() is not database2.connection()
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_connection_context_with_raw_connection(database_url):
@@ -961,16 +976,59 @@ async def test_database_url_interface(database_url):
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_concurrent_access_on_single_connection(database_url):
-    database_url = DatabaseURL(database_url)
-    if database_url.dialect != "postgresql":
-        pytest.skip("Test requires `pg_sleep()`")
-
     async with Database(database_url, force_rollback=True) as database:
 
         async def db_lookup():
-            await database.fetch_one("SELECT pg_sleep(1)")
+            await database.fetch_one("SELECT 1 AS value")
+
+        await asyncio.gather(
+            db_lookup(),
+            db_lookup(),
+        )
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_concurrent_transactions_on_single_connection(database_url: str):
+    async with Database(database_url) as database:
+
+        @database.transaction()
+        async def db_lookup():
+            await database.fetch_one(query="SELECT 1 AS value")
+
+        await asyncio.gather(
+            db_lookup(),
+            db_lookup(),
+        )
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_concurrent_tasks_on_single_connection(database_url: str):
+    async with Database(database_url) as database:
+
+        async def db_lookup():
+            await database.fetch_one(query="SELECT 1 AS value")
+
+        await asyncio.gather(
+            asyncio.create_task(db_lookup()),
+            asyncio.create_task(db_lookup()),
+        )
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_concurrent_task_transactions_on_single_connection(database_url: str):
+    async with Database(database_url) as database:
+
+        @database.transaction()
+        async def db_lookup():
+            await database.fetch_one(query="SELECT 1 AS value")
 
-        await asyncio.gather(db_lookup(), db_lookup())
+        await asyncio.gather(
+            asyncio.create_task(db_lookup()),
+            asyncio.create_task(db_lookup()),
+        )
 
 
 @pytest.mark.parametrize("database_url", DATABASE_URLS)