test: add new tests for upcoming contextvar inheritance/isolation and weakref cleanup

Author: zevisert

tests/test_databases.py

@@ -5,8 +5,9 @@
 import itertools
 import os
 import re
+import gc
 from unittest.mock import MagicMock, patch
-
+from typing import MutableMapping
 import pytest
 import sqlalchemy
 
@@ -478,6 +479,212 @@ async def test_transaction_commit(database_url):
             assert len(results) == 1
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_transaction_context_child_task_interaction(database_url):
+    """
+    Ensure that child tasks may influence inherited transactions.
+    """
+    # This is an practical example of the next test.
+    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():
+                # Chage the note from a child task
+                await database.execute(
+                    notes.update().where(notes.c.id == 1).values(text="test")
+                )
+
+            await asyncio.create_task(run_update_from_child_task())
+
+            # 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_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_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_connection_context_cleanup_contextmanager(database_url):
+    """
+    Ensure that contextvar connections are not persisted unecessarily.
+    """
+    from databases.core import _ACTIVE_CONNECTIONS
+
+    assert _ACTIVE_CONNECTIONS.get() is None
+
+    async with Database(database_url) as database:
+        # .connect is lazy, it doesn't create a Connection, but .connection does
+        connection = database.connection()
+
+        open_connections = _ACTIVE_CONNECTIONS.get()
+        assert isinstance(open_connections, MutableMapping)
+        assert open_connections.get(database) is connection
+
+    # Context manager closes, open_connections is cleaned up
+    open_connections = _ACTIVE_CONNECTIONS.get()
+    assert isinstance(open_connections, MutableMapping)
+    assert open_connections.get(database, None) is None
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_connection_context_cleanup_garbagecollector(database_url):
+    """
+    Ensure that contextvar connections are not persisted unecessarily, even
+    if exit handlers are not called.
+    """
+    from databases.core import _ACTIVE_CONNECTIONS
+
+    assert _ACTIVE_CONNECTIONS.get() is None
+
+    database = Database(database_url)
+    await database.connect()
+    connection = database.connection()
+
+    # Should be tracking the connection
+    open_connections = _ACTIVE_CONNECTIONS.get()
+    assert isinstance(open_connections, MutableMapping)
+    assert open_connections.get(database) is connection
+
+    # neither .disconnect nor .__aexit__ are called before deleting the reference
+    del database
+    gc.collect()
+
+    # Should have dropped reference to connection, even without proper cleanup
+    open_connections = _ACTIVE_CONNECTIONS.get()
+    assert isinstance(open_connections, MutableMapping)
+    assert len(open_connections) == 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 databases.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 databases.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()
+
+        # TODO(zevisert,review): Could skip instead of using the logic below
+        # 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
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_transaction_commit_serializable(database_url):