Auto-recreate connection on loop change (#2098)

* Auto-recreate connection on loop change

* Fix test

Author: abondar

docs/connections.rst

@@ -162,6 +162,54 @@ providing isolation between different contexts (useful for testing).
     conn3 = Tortoise.get_connection("default")
     assert conn is not conn3
 
+Event Loop Handling
+===================
+
+Some database drivers (asyncpg, aiomysql) bind their connection pools to the event loop
+that created them. If the loop changes -- for example, when using function-scoped pytest
+fixtures or Starlette's ``TestClient`` -- the old pool becomes unusable.
+
+Tortoise handles this automatically: when ``ConnectionHandler.get()`` detects that the
+current event loop differs from the one the connection was created on, it transparently
+creates a fresh connection.
+
+**In production**, a loop change usually indicates a bug (e.g., mixing sync/async code).
+A ``TortoiseLoopSwitchWarning`` is emitted so you can investigate:
+
+.. code-block:: python
+
+    import warnings
+    from tortoise.warnings import TortoiseLoopSwitchWarning
+
+    # Suppress if you know what you're doing
+    warnings.filterwarnings("ignore", category=TortoiseLoopSwitchWarning)
+
+**In tests**, ``tortoise_test_context()`` suppresses this warning automatically.
+No special configuration needed.
+
+.. list-table:: Backend Loop Binding
+   :header-rows: 1
+   :widths: 30 20 50
+
+   * - Backend
+     - Bound?
+     - Notes
+   * - asyncpg
+     - Yes
+     - Pool stores loop at creation time
+   * - aiomysql/asyncmy
+     - Yes
+     - Pool stores loop at creation time
+   * - psycopg
+     - No
+     - Uses running loop per-operation
+   * - aiosqlite
+     - Partial
+     - Grabs loop per-operation, not at creation
+   * - asyncodbc (MSSQL/Oracle)
+     - No
+     - Per-operation loop resolution
+
 API Reference
 =============
 

docs/contrib/unittest.rst

@@ -123,6 +123,31 @@ For tests that require multiple database connections:
             await ctx.generate_schemas()
             yield ctx
 
+Event Loop Isolation
+====================
+
+Some backends (asyncpg, aiomysql) bind connection pools to the event loop that created
+them. ``tortoise_test_context()`` handles this transparently -- if the event loop changes
+between tests, connections are automatically recreated.
+
+This means you **don't** need ``loop_scope="session"`` or any special pytest-asyncio
+configuration. The simplest setup works:
+
+.. code-block:: toml
+
+    # pyproject.toml -- no loop_scope overrides needed
+    [tool.pytest.ini_options]
+    asyncio_mode = "auto"
+
+If you use ``TortoiseContext`` directly (without ``tortoise_test_context``), you may see
+a ``TortoiseLoopSwitchWarning`` when the loop changes. Suppress it with:
+
+.. code-block:: python
+
+    import warnings
+    from tortoise.warnings import TortoiseLoopSwitchWarning
+    warnings.filterwarnings("ignore", category=TortoiseLoopSwitchWarning)
+
 Testing Database Capabilities
 =============================
 

tests/test_connection.py

@@ -1,9 +1,12 @@
+import asyncio
+import warnings
 from unittest.mock import AsyncMock, Mock, PropertyMock, call, patch
 
 import pytest
 
 from tortoise import BaseDBAsyncClient, ConfigurationError
 from tortoise.connection import ConnectionHandler
+from tortoise.warnings import TortoiseLoopSwitchWarning
 
 
 @pytest.fixture
@@ -206,9 +209,10 @@ def test_create_connection_db_info_not_str(
 
 
 def test_get_alias_present(conn_handler):
-    conn_handler._storage = {"default": "some_connection"}
+    mock_conn = Mock(_check_loop=Mock(return_value=True))
+    conn_handler._storage = {"default": mock_conn}
     ret_val = conn_handler.get("default")
-    assert ret_val == "some_connection"
+    assert ret_val is mock_conn
 
 
 @patch("tortoise.connection.ConnectionHandler._create_connection")
@@ -246,10 +250,12 @@ def test_reset(conn_handler):
 
 @patch("tortoise.connection.ConnectionHandler.db_config", new_callable=PropertyMock)
 def test_all(mocked_db_config, conn_handler):
-    conn_handler._storage = {"default": "some_conn", "other": "some_other_conn"}
+    mock_conn_1 = Mock(_check_loop=Mock(return_value=True))
+    mock_conn_2 = Mock(_check_loop=Mock(return_value=True))
+    conn_handler._storage = {"default": mock_conn_1, "other": mock_conn_2}
     mocked_db_config.return_value = {"default": {}, "other": {}}
     ret_val = conn_handler.all()
-    assert set(ret_val) == {"some_conn", "some_other_conn"}
+    assert set(ret_val) == {mock_conn_1, mock_conn_2}
 
 
 @pytest.mark.asyncio
@@ -282,3 +288,50 @@ async def test_close_all_without_discard(mocked_db_config, conn_handler):
     conn_1.close.assert_awaited_once()
     conn_2.close.assert_awaited_once()
     assert conn_handler._storage == {"default": conn_1, "other": conn_2}
+
+
+# --- Event loop validation tests ---
+
+
+@pytest.mark.asyncio
+async def test_check_loop_returns_true_when_not_bound():
+    client = Mock(spec=BaseDBAsyncClient)
+    client._bound_loop = None
+    client._check_loop = BaseDBAsyncClient._check_loop.__get__(client)
+    assert client._check_loop() is True
+
+
+@pytest.mark.asyncio
+async def test_check_loop_returns_true_on_same_loop():
+    client = Mock(spec=BaseDBAsyncClient)
+    client._bound_loop = asyncio.get_running_loop()
+    client._check_loop = BaseDBAsyncClient._check_loop.__get__(client)
+    assert client._check_loop() is True
+
+
+@pytest.mark.asyncio
+async def test_check_loop_returns_false_on_different_loop():
+    client = Mock(spec=BaseDBAsyncClient)
+    client._bound_loop = asyncio.new_event_loop()
+    client._check_loop = BaseDBAsyncClient._check_loop.__get__(client)
+    assert client._check_loop() is False
+
+
+@patch("tortoise.connection.ConnectionHandler._create_connection")
+def test_get_reconnects_on_loop_change(mocked_create_connection, conn_handler):
+    """When _check_loop() returns False, get() should warn and create a new connection."""
+    stale_conn = Mock(_check_loop=Mock(return_value=False))
+    fresh_conn = Mock(_check_loop=Mock(return_value=True))
+    mocked_create_connection.return_value = fresh_conn
+    conn_handler._storage = {"default": stale_conn}
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        ret_val = conn_handler.get("default")
+
+    assert ret_val is fresh_conn
+    assert conn_handler._storage["default"] is fresh_conn
+    mocked_create_connection.assert_called_once_with("default")
+    loop_warnings = [x for x in w if issubclass(x.category, TortoiseLoopSwitchWarning)]
+    assert len(loop_warnings) == 1
+    assert "different event loop" in str(loop_warnings[0].message)

tests/test_query_api.py

@@ -278,6 +278,10 @@ async def test_execute_pypika_explicit_connection_with_multiple_configured() ->
 
     class DummyClient:
         query_class = type("QueryClass", (), {"SQL_CONTEXT": None})
+        _bound_loop = None
+
+        def _check_loop(self) -> bool:
+            return True
 
         async def execute_query_dict_with_affected(self, query, values=None):
             return [], 0

tortoise/backends/asyncpg/client.py

@@ -59,6 +59,7 @@ async def create_connection(self, with_db: bool) -> None:
         }
         try:
             self._pool = await self.create_pool(password=self.password, **self._template)
+            await self._post_connect()
             self.log.debug("Created connection pool %s with params: %s", self._pool, self._template)
         except asyncpg.InvalidCatalogNameError as ex:
             msg = "Can't establish connection to "

tortoise/backends/base/client.py

@@ -118,6 +118,7 @@ class BaseDBAsyncClient(abc.ABC):
     _connection: Any
     _parent: BaseDBAsyncClient
     _pool: Any
+    _bound_loop: asyncio.AbstractEventLoop | None = None
     connection_name: str
     query_class: type[Query] = Query
     executor_class: type[BaseExecutor] = BaseExecutor
@@ -129,6 +130,20 @@ def __init__(self, connection_name: str, fetch_inserted: bool = True, **kwargs:
         self.connection_name = connection_name
         self.fetch_inserted = fetch_inserted
 
+    def _check_loop(self) -> bool:
+        """Check if the current event loop matches the one this client was created on."""
+        try:
+            current = asyncio.get_running_loop()
+        except RuntimeError:
+            return True  # No running loop — can't validate
+        if self._bound_loop is None:
+            return True  # Not yet bound (pool not created yet)
+        return self._bound_loop is current
+
+    async def _post_connect(self) -> None:
+        """Called after pool/connection is created. Records the bound loop."""
+        self._bound_loop = asyncio.get_running_loop()
+
     async def create_connection(self, with_db: bool) -> None:
         """
         Establish a DB connection.

tortoise/backends/mysql/client.py

@@ -140,6 +140,7 @@ async def create_connection(self, with_db: bool) -> None:
                             hours = timezone.now().utcoffset().seconds / 3600  # type: ignore
                             tz = f"{int(hours):+d}:{int((hours % 1) * 60):02d}"
                             await cursor.execute(f"SET time_zone='{tz}';")
+            await self._post_connect()
             self.log.debug("Created connection %s pool with params: %s", self._pool, self._template)
         except errors.OperationalError:
             raise DBConnectionError(f"Can't connect to MySQL server: {self._template}")

tortoise/backends/odbc/client.py

@@ -91,6 +91,7 @@ async def create_connection(self, with_db: bool) -> None:
             self._pool = await asyncodbc.create_pool(
                 **self._template,
             )
+            await self._post_connect()
             self.log.debug("Created connection %s pool with params: %s", self._pool, self._template)
         except pyodbc.InterfaceError:
             raise DBConnectionError(f"Can't establish connection to database {self.database}")

tortoise/backends/psycopg/client.py

@@ -104,6 +104,7 @@ async def create_connection(self, with_db: bool) -> None:
             # Immediately test the connection because the test suite expects it to check if the
             # connection is valid.
             await self._pool.open(wait=True, timeout=extra["timeout"])
+            await self._post_connect()
             self.log.debug("Created connection pool %s with params: %s", self._pool, self._template)
         except (psycopg.errors.InvalidCatalogName, psycopg_pool.PoolTimeout):
             raise exceptions.DBConnectionError(

tortoise/backends/sqlite/client.py

@@ -84,6 +84,7 @@ async def create_connection(self, with_db: bool) -> None:
             for pragma, val in self.pragmas.items():
                 cursor = await self._connection.execute(f"PRAGMA {pragma}={val}")
                 await cursor.close()
+            await self._post_connect()
             self.log.debug(
                 "Created connection %s with params: filename=%s %s",
                 self._connection,