Merge remote-tracking branch 'madkinsz/example/instance-safe' into fix-transaction-contextvar

Author: zevisert

databases/core.py

@@ -3,9 +3,7 @@
 import functools
 import logging
 import typing
-from contextvars import ContextVar
 from types import TracebackType
-from typing import Dict, Tuple
 from urllib.parse import SplitResult, parse_qsl, unquote, urlsplit
 
 from sqlalchemy import text
@@ -35,21 +33,6 @@
 
 logger = logging.getLogger("databases")
 
-# Connections are stored as task-local state, but care must be taken to ensure
-# that two database instances in the same task do not overwrite each other's connections.
-# For this reason, the dict key comprises the database instance and the current task.
-_connection_contextmap: ContextVar[
-    Dict[Tuple["Database", asyncio.Task], "Connection"]
-] = ContextVar("databases:Connection")
-
-
-def _get_connection_contextmap() -> Dict[Tuple["Database", asyncio.Task], "Connection"]:
-    connections = _connection_contextmap.get(None)
-    if connections is None:
-        connections = {}
-        _connection_contextmap.set(connections)
-    return connections
-
 
 class Database:
     SUPPORTED_BACKENDS = {
@@ -79,6 +62,9 @@ def __init__(
         assert issubclass(backend_cls, DatabaseBackend)
         self._backend = backend_cls(self.url, **self.options)
 
+        # Connections are stored per asyncio task
+        self._connections: typing.Dict[typing.Optional[asyncio.Task], Connection] = {}
+
         # When `force_rollback=True` is used, we use a single global
         # connection, within a transaction that always rolls back.
         self._global_connection: typing.Optional[Connection] = None
@@ -126,11 +112,10 @@ async def disconnect(self) -> None:
             self._global_transaction = None
             self._global_connection = None
         else:
-            task = asyncio.current_task()
-            assert task is not None, "Not running in an asyncio task"
-            connections = _get_connection_contextmap()
-            if (self, task) in connections:
-                del connections[self, task]
+            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(
@@ -204,13 +189,12 @@ def connection(self) -> "Connection":
         if self._global_connection is not None:
             return self._global_connection
 
-        task = asyncio.current_task()
-        assert task is not None, "Not running in an asyncio task"
-        connections = _get_connection_contextmap()
-        if (self, task) not in connections:
-            connections[self, task] = Connection(self._backend)
+        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 connections[self, task]
+        return self._connections[current_task]
 
     def transaction(
         self, *, force_rollback: bool = False, **kwargs: typing.Any
@@ -351,19 +335,6 @@ def _build_query(
 
 _CallableType = typing.TypeVar("_CallableType", bound=typing.Callable)
 
-_transaction_contextmap: ContextVar[
-    Dict["Transaction", TransactionBackend]
-] = ContextVar("databases:Transactions")
-
-
-def _get_transaction_contextmap() -> Dict["Transaction", TransactionBackend]:
-    transactions = _transaction_contextmap.get(None)
-    if transactions is None:
-        transactions = {}
-        _transaction_contextmap.set(transactions)
-
-    return transactions
-
 
 class Transaction:
     def __init__(
@@ -376,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()`
@@ -417,9 +393,10 @@ async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
 
     async def start(self) -> "Transaction":
         connection = self._connection_callable()
+        current_task = asyncio.current_task()
+        assert current_task is not None, "No currently running task"
         transaction = connection._connection.transaction()
-        transactions = _get_transaction_contextmap()
-        transactions[self] = transaction
+        self._transactions[current_task] = transaction
         async with connection._transaction_lock:
             is_root = not connection._transaction_stack
             await connection.__aenter__()
@@ -429,27 +406,27 @@ async def start(self) -> "Transaction":
 
     async def commit(self) -> None:
         connection = self._connection_callable()
-        transactions = _get_transaction_contextmap()
-        transaction = transactions.get(self, None)
+        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 transactions[self]
+            del self._transactions[current_task]
 
     async def rollback(self) -> None:
         connection = self._connection_callable()
-        transactions = _get_transaction_contextmap()
-        transaction = transactions.get(self, None)
+        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 transactions[self]
+            del self._transactions[current_task]
 
 
 class _EmptyNetloc(str):