diff --git a/databasez/__init__.py b/databasez/__init__.py index 20e9715..8e8228b 100644 --- a/databasez/__init__.py +++ b/databasez/__init__.py @@ -1,5 +1,5 @@ from databasez.core import Database, DatabaseURL -__version__ = "0.4.0" +__version__ = "0.5.0" __all__ = ["Database", "DatabaseURL"] diff --git a/databasez/core.py b/databasez/core.py index 30a5b12..db32b1c 100644 --- a/databasez/core.py +++ b/databasez/core.py @@ -3,16 +3,16 @@ import functools import logging import typing +import weakref from contextvars import ContextVar from types import TracebackType from urllib.parse import SplitResult, parse_qsl, unquote, urlencode, urlsplit from sqlalchemy import text from sqlalchemy.sql import ClauseElement -from sqlalchemy.util._concurrency_py3k import greenlet_spawn from databasez.importer import import_from_string -from databasez.interfaces import DatabaseBackend, Record +from databasez.interfaces import DatabaseBackend, Record, TransactionBackend if typing.TYPE_CHECKING: from databasez.types import DictAny @@ -35,6 +35,11 @@ logger = logging.getLogger("databasez") +ACTIVE_TRANSACTIONS: ContextVar[ + typing.Optional["weakref.WeakKeyDictionary['Transaction', 'TransactionBackend']"] +] = ContextVar("databasez:active_transactions", default=None) + + class Database: """ An abstraction on the top of the EncodeORM databases.Database object. @@ -72,6 +77,7 @@ class Database: } DIRECT_URL_SCHEME = {"sqlite"} MANDATORY_FIELDS = ["host", "port", "user", "database"] + _connection_map: "weakref.WeakKeyDictionary[asyncio.Task, 'Connection']" def __init__( self, @@ -92,6 +98,7 @@ def __init__( self.url = DatabaseURL(_url) # type: ignore self.options = options self.is_connected = False + self._connection_map = weakref.WeakKeyDictionary() self._force_rollback = force_rollback @@ -100,9 +107,6 @@ 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") - # 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 @@ -164,6 +168,30 @@ def _build_url_for_direct_url_scheme(self, scheme: str, database: str) -> str: """ return f"{scheme}:///{database}" + @property + def _current_task(self) -> asyncio.Task: + task = asyncio.current_task() + if not task: + raise RuntimeError("No currently active asyncio.Task found") + return task + + @property + def _connection(self) -> typing.Optional["Connection"]: + return self._connection_map.get(self._current_task) + + @_connection.setter + def _connection( + self, connection: typing.Optional["Connection"] + ) -> typing.Optional["Connection"]: + task = self._current_task + + if connection is None: + self._connection_map.pop(task, None) + else: + self._connection_map[task] = connection + + return self._connection + async def connect(self) -> None: """ Establish the connection pool. @@ -180,7 +208,7 @@ async def connect(self) -> None: assert self._global_connection is None assert self._global_transaction is None - self._global_connection = Connection(self._backend) + self._global_connection = Connection(self, self._backend) self._global_transaction = self._global_connection.transaction(force_rollback=True) await self._global_transaction.__aenter__() @@ -202,7 +230,7 @@ async def disconnect(self) -> None: self._global_transaction = None self._global_connection = None else: - self._connection_context = ContextVar("connection_context") + self._connection = None await self._backend.disconnect() logger.info( @@ -274,12 +302,9 @@ def connection(self) -> "Connection": if self._global_connection is not None: return self._global_connection - try: - return self._connection_context.get() # type: ignore - except LookupError: - connection = Connection(self._backend) - self._connection_context.set(connection) - return connection + if not self._connection: + self._connection = Connection(self, self._backend) + return self._connection def transaction(self, *, force_rollback: bool = False, **kwargs: typing.Any) -> "Transaction": return Transaction(self.connection, force_rollback=force_rollback, **kwargs) @@ -300,7 +325,8 @@ def _get_backend(self) -> str: class Connection: - def __init__(self, backend: DatabaseBackend) -> None: + def __init__(self, database: Database, backend: DatabaseBackend) -> None: + self._database = database self._backend = backend self._connection_lock = asyncio.Lock() @@ -334,6 +360,7 @@ async def __aexit__( self._connection_counter -= 1 if self._connection_counter == 0: await self._connection.release() + self._database._connection = None async def fetch_all( self, @@ -398,11 +425,6 @@ def connection_callable() -> Connection: def raw_connection(self) -> typing.Any: return self._connection.raw_connection - async def run_sync( - self, fn: typing.Callable[..., typing.Any], *arg: typing.Any, **kw: typing.Any - ) -> typing.Any: - return await greenlet_spawn(fn, self._connection.raw_connection, *arg, **kw) - @staticmethod def _build_query( query: typing.Union[ClauseElement, str], values: typing.Optional[dict] = None @@ -431,6 +453,37 @@ def __init__( self._force_rollback = force_rollback self._extra_options = kwargs + @property + def _connection(self) -> "Connection": + # Returns the same connection if called multiple times + return self._connection_callable() + + @property + def _transaction(self) -> typing.Optional["TransactionBackend"]: + transactions = ACTIVE_TRANSACTIONS.get() + if transactions is None: + return None + + return transactions.get(self, None) + + @_transaction.setter + def _transaction( + self, transaction: typing.Optional["TransactionBackend"] + ) -> typing.Optional["TransactionBackend"]: + transactions = ACTIVE_TRANSACTIONS.get() + if transactions is None: + transactions = weakref.WeakKeyDictionary() + else: + transactions = transactions.copy() + + if transaction is None: + transactions.pop(self, None) + else: + transactions[self] = transaction + + ACTIVE_TRANSACTIONS.set(transactions) + return transactions.get(self, None) + async def __aenter__(self) -> "Transaction": """ Called when entering `async with database.transaction()` @@ -471,7 +524,6 @@ 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: @@ -485,15 +537,19 @@ async def commit(self) -> None: async with self._connection._transaction_lock: assert self._connection._transaction_stack[-1] is self self._connection._transaction_stack.pop() + assert self._transaction is not None await self._transaction.commit() await self._connection.__aexit__() + self._transaction = None async def rollback(self) -> None: async with self._connection._transaction_lock: assert self._connection._transaction_stack[-1] is self self._connection._transaction_stack.pop() + assert self._transaction is not None await self._transaction.rollback() await self._connection.__aexit__() + self._transaction = None class _EmptyNetloc(str): diff --git a/docs/connections-and-transactions.md b/docs/connections-and-transactions.md index 8745081..7973771 100644 --- a/docs/connections-and-transactions.md +++ b/docs/connections-and-transactions.md @@ -37,14 +37,14 @@ to connect to the database. ## Connecting and disconnecting -You can control the database connect/disconnect, by using it as a async context manager. +You can control the database connection, by using it as a async context manager. ```python async with Database(DATABASE_URL) as database: ... ``` -Or by using explicit connection and disconnection: +Or by using explicit `.connect()` and `disconnect()`: ```python database = Database(DATABASE_URL) @@ -246,11 +246,54 @@ async def create_users(request): ... ``` -Transaction blocks are managed as task-local state. Nested transactions -are fully supported, and are implemented using database savepoints. +The state of a transaction is liked to the connection used in the currently executing async task. +If you would like to influence an active transaction from another task, the connection must be +shared: Transaction isolation-level can be specified if the driver backend supports that: +```python +async def add_excitement(connnection: databases.core.Connection, id: int): + await connection.execute( + "UPDATE notes SET text = CONCAT(text, '!!!') WHERE id = :id", + {"id": id} + ) + + +async with Database(database_url) as database: + async with database.transaction(): + # This note won't exist until the transaction closes... + await database.execute( + "INSERT INTO notes(id, text) values (1, 'databases is cool')" + ) + # ...but child tasks can use this connection now! + await asyncio.create_task(add_excitement(database.connection(), id=1)) + + await database.fetch_val("SELECT text FROM notes WHERE id=1") + # ^ returns: "databases is cool!!!" +``` + +Nested transactions are fully supported, and are implemented using database savepoints: + +```python +async with databases.Database(database_url) as db: + async with db.transaction() as outer: + # Do something in the outer transaction + ... + + # Suppress to prevent influence on the outer transaction + with contextlib.suppress(ValueError): + async with db.transaction(): + # Do something in the inner transaction + ... + + raise ValueError('Abort the inner transaction') + + # Observe the results of the outer transaction, + # without effects from the inner transaction. + await db.fetch_all('SELECT * FROM ...') +``` + ```python async with database.transaction(isolation="serializable"): ... diff --git a/docs/release-notes.md b/docs/release-notes.md index ac25f36..806622c 100644 --- a/docs/release-notes.md +++ b/docs/release-notes.md @@ -1,5 +1,13 @@ # Release Notes +## 0.5.0 + +### Fixed + +- Patch done in the core of Databases fixing the concurrent usage of connections and transactions. +This patch also affects databases. [#PR 546](https://github.com/encode/databases/pull/546) by [@zevisert](https://github.com/zevisert). +We thank [@zevisert](https://github.com/zevisert) for the fix done in the original project that also affect Databasez. + ## 0.4.0 ### Changed diff --git a/docs_src/testclient/tests.py b/docs_src/testclient/tests.py index c56ce1b..0826a11 100644 --- a/docs_src/testclient/tests.py +++ b/docs_src/testclient/tests.py @@ -7,7 +7,8 @@ import pytest import saffier from databasez.testclient import DatabaseTestClient -from saffier import fields +from saffier.db.models import fields + from tests.settings import DATABASE_URL database = DatabaseTestClient(DATABASE_URL, drop_database=True) diff --git a/pyproject.toml b/pyproject.toml index 61b1029..0cae168 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ "Topic :: Internet :: WWW/HTTP :: HTTP Servers", "Topic :: Internet :: WWW/HTTP", ] -dependencies = ["nest_asyncio>=1.5.6,<2.0.0", "sqlalchemy>=2.0.16,<2.1"] +dependencies = ["nest_asyncio>=1.5.6,<2.0.0", "sqlalchemy>=2.0.19,<2.1"] keywords = [ "mysql", "postgres", @@ -78,6 +78,7 @@ dev = [ "aiosqlite>=0.18.0,<0.20.0", "asyncpg>=0.27.0,<0.30.0", "aioodbc>=0.4.0,<0.5.0", + "ipdb>=0.13.13", "pre-commit>=2.17.0,<4.0.0", "psycopg2-binary>=2.9.6,<3.0.0", "pymysql>=1.0.3,<2.0.0", diff --git a/tests/test_databases.py b/tests/test_databases.py index 36868a0..23c6b61 100644 --- a/tests/test_databases.py +++ b/tests/test_databases.py @@ -2,7 +2,9 @@ import datetime import decimal import functools +import gc import os +from typing import MutableMapping from unittest.mock import MagicMock, patch from urllib.parse import parse_qsl, urlsplit @@ -1354,3 +1356,246 @@ async def test_mapping_property_interface(database_url): list_result = await database.fetch_all(query=query) assert list_result[0]._mapping["text"] == "example1" assert list_result[0]._mapping["completed"] is True + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_child_task_inheritance(database_url): + """ + Ensure that transactions are inherited by child tasks. + """ + async with Database(database_url) as database: + + async def check_transaction(transaction, active_transaction): + # Should have inherited the same transaction backend from the parent task + assert transaction._transaction is active_transaction + + async with database.transaction() as transaction: + await asyncio.create_task(check_transaction(transaction, transaction._transaction)) + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_child_task_inheritance_example(database_url): + """ + Ensure that child tasks may influence inherited transactions. + """ + # This is an practical example of the above test. + db = Database(database_url) + if db.url.dialect == "mssql": + return + + async with Database(database_url) as database: + async with database.transaction(): + # Create a note + await database.execute(notes.insert().values(id=1, text="setup", completed=True)) + + # Change the note from the same task + await database.execute(notes.update().where(notes.c.id == 1).values(text="prior")) + + # Confirm the change + result = await database.fetch_one(notes.select().where(notes.c.id == 1)) + assert result.text == "prior" + + async def run_update_from_child_task(connection): + # Change the note from a child task + await connection.execute(notes.update().where(notes.c.id == 1).values(text="test")) + + await asyncio.create_task(run_update_from_child_task(database.connection())) + + # Confirm the child's change + result = await database.fetch_one(notes.select().where(notes.c.id == 1)) + assert result.text == "test" + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_sibling_task_isolation(database_url): + """ + Ensure that transactions are isolated between sibling tasks. + """ + start = asyncio.Event() + end = asyncio.Event() + + async with Database(database_url) as database: + + async def check_transaction(transaction): + await start.wait() + # Parent task is now in a transaction, we should not + # see its transaction backend since this task was + # _started_ in a context where no transaction was active. + assert transaction._transaction is None + end.set() + + transaction = database.transaction() + assert transaction._transaction is None + task = asyncio.create_task(check_transaction(transaction)) + + async with transaction: + start.set() + assert transaction._transaction is not None + await end.wait() + + # Cleanup for "Task not awaited" warning + await task + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_sibling_task_isolation_example(database_url): + """ + Ensure that transactions are running in sibling tasks are isolated from eachother. + """ + # This is an practical example of the above test. + db = Database(database_url) + if db.url.dialect == "mssql": + return + setup = asyncio.Event() + done = asyncio.Event() + + async def tx1(connection): + async with connection.transaction(): + await db.execute(notes.insert(), values={"id": 1, "text": "tx1", "completed": False}) + setup.set() + await done.wait() + + async def tx2(connection): + async with connection.transaction(): + await setup.wait() + result = await db.fetch_all(notes.select()) + assert result == [], result + done.set() + + async with Database(database_url) as db: + await asyncio.gather(tx1(db), tx2(db)) + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_connection_cleanup_contextmanager(database_url): + """ + Ensure that task connections are not persisted unecessarily. + """ + + ready = asyncio.Event() + done = asyncio.Event() + + async def check_child_connection(database: Database): + async with database.connection(): + ready.set() + await done.wait() + + async with Database(database_url) as database: + # Should have a connection in this task + # .connect is lazy, it doesn't create a Connection, but .connection does + connection = database.connection() + assert isinstance(database._connection_map, MutableMapping) + assert database._connection_map.get(asyncio.current_task()) is connection + + # Create a child task and see if it registers a connection + task = asyncio.create_task(check_child_connection(database)) + await ready.wait() + assert database._connection_map.get(task) is not None + assert database._connection_map.get(task) is not connection + + # Let the child task finish, and see if it cleaned up + done.set() + await task + # This is normal exit logic cleanup, the WeakKeyDictionary + # shouldn't have cleaned up yet since the task is still referenced + assert task not in database._connection_map + + # Context manager closes, all open connections are removed + assert isinstance(database._connection_map, MutableMapping) + assert len(database._connection_map) == 0 + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_connection_cleanup_garbagecollector(database_url): + """ + Ensure that connections for tasks are not persisted unecessarily, even + if exit handlers are not called. + """ + database = Database(database_url) + await database.connect() + + created = asyncio.Event() + + async def check_child_connection(database: Database): + # neither .disconnect nor .__aexit__ are called before deleting this task + database.connection() + created.set() + + task = asyncio.create_task(check_child_connection(database)) + await created.wait() + assert task in database._connection_map + await task + del task + gc.collect() + + # Should not have a connection for the task anymore + assert len(database._connection_map) == 0 + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_cleanup_contextmanager(database_url): + """ + Ensure that contextvar transactions are not persisted unecessarily. + """ + from databasez.core import ACTIVE_TRANSACTIONS + + assert ACTIVE_TRANSACTIONS.get() is None + + async with Database(database_url) as database: + async with database.transaction() as transaction: + open_transactions = ACTIVE_TRANSACTIONS.get() + assert isinstance(open_transactions, MutableMapping) + assert open_transactions.get(transaction) is transaction._transaction + + # Context manager closes, open_transactions is cleaned up + open_transactions = ACTIVE_TRANSACTIONS.get() + assert isinstance(open_transactions, MutableMapping) + assert open_transactions.get(transaction, None) is None + + +@pytest.mark.parametrize("database_url", DATABASE_URLS) +@async_adapter +async def test_transaction_context_cleanup_garbagecollector(database_url): + """ + Ensure that contextvar transactions are not persisted unecessarily, even + if exit handlers are not called. + This test should be an XFAIL, but cannot be due to the way that is hangs + during teardown. + """ + from databasez.core import ACTIVE_TRANSACTIONS + + assert ACTIVE_TRANSACTIONS.get() is None + + async with Database(database_url) as database: + transaction = database.transaction() + await transaction.start() + + # Should be tracking the transaction + open_transactions = ACTIVE_TRANSACTIONS.get() + assert isinstance(open_transactions, MutableMapping) + assert open_transactions.get(transaction) is transaction._transaction + + # neither .commit, .rollback, nor .__aexit__ are called + del transaction + gc.collect() + + # A strong reference to the transaction is kept alive by the connection's + # ._transaction_stack, so it is still be tracked at this point. + assert len(open_transactions) == 1 + + # If that were magically cleared, the transaction would be cleaned up, + # but as it stands this always causes a hang during teardown at + # `Database(...).disconnect()` if the transaction is not closed. + transaction = database.connection()._transaction_stack[-1] + await transaction.rollback() + del transaction + + # Now with the transaction rolled-back, it should be cleaned up. + assert len(open_transactions) == 0