Repository: PrefectHQ/prefect

Issue:
Connection leak in the test suite (from `pluggy>=1.4.0` warnings)
```
tests/server/api/test_block_documents.py::TestCreateBlockDocument::test_create_block_document_with_nonstandard_characters[my-block] - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
Plugin: unraisableexception, Hook: pytest_runtest_setup
PytestUnraisableExceptionWarning: Exception ignored in: <function Connection.__del__ at 0x7fe03fc0d4e0>
```

This has been popping up in random tests since last week, and it is because `pluggy` is giving us a new warning: https://github.com/pytest-dev/pluggy/compare/1.3.0...1.4.0

```
- `#463 <https://github.com/pytest-dev/pluggy/issues/463>`_: A warning :class:`~pluggy.PluggyTeardownRaisedWarning` is now issued when an old-style hookwrapper raises an exception during teardown.
  See the warning documentation for more details.
```

I'll be temporarily suppressing that warning, and this issue can be for tracking down and fixing that leaked connection.

It appears that the error is always in python 3.8 and always with the postgres DB backend, which may be a hint.

Examples of runs that failed with this error:

```
Job ID: 21162420883
Job Name: python:3.8, postgres:14, --exclude-services
Started at: 2024-02-02T17:01:33Z
Link: https://github.com/PrefectHQ/prefect/actions/runs/7759102494/job/21162420883
2024-02-02T17:11:01.8665082Z [31mFAILED[0m tests/server/test_app.py::[1mtest_app_exposes_ui_settings[0m - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
```

```
Job ID: 21156157922
Job Name: python:3.8, postgres:13, --exclude-services
Started at: 2024-02-02T14:18:43Z
Link: https://github.com/PrefectHQ/prefect/actions/runs/7757198583/job/21156157922
2024-02-02T14:28:17.5462675Z [31mFAILED[0m tests/server/test_app.py::[1mtest_app_exposes_ui_settings[0m - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
```

```
Job ID: 21112206170
Job Name: pydantic v1, python:3.8, postgres:14, --exclude-services
Started at: 2024-02-01T14:17:16Z
Link: https://github.com/PrefectHQ/prefect/actions/runs/7742635501/job/21112206170
2024-02-01T14:26:38.8923204Z [31mFAILED[0m tests/server/test_app.py::[1mtest_app_exposes_ui_settings[0m - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
```

```
Job ID: 20991318569
Job Name: python:3.8, postgres:13, --exclude-services
Started at: 2024-01-29T21:25:03Z
Link: https://github.com/PrefectHQ/prefect/actions/runs/7702627286/job/20991318569
2024-01-29T21:34:31.4565524Z [31mFAILED[0m tests/server/test_app.py::[1mtest_app_exposes_ui_settings[0m - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
```

```
Job ID: 20874540746
Job Name: pydantic v1, python:3.8, postgres:14, --exclude-services
Started at: 2024-01-25T19:23:40Z
Link: https://github.com/PrefectHQ/prefect/actions/runs/7659435084/job/20874540746
2024-01-25T19:33:08.7407781Z [31mFAILED[0m tests/server/test_app.py::[1mtest_app_exposes_ui_settings[0m - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
```



Discussion:
Thinking now that @abrookins may have fixed this, going to try again with pluggy unpinned
https://github.com/PrefectHQ/prefect/actions/runs/7817232549/job/21324622376?pr=11916

Some better output from a recent failure:

```
ERROR tests/input/test_run_input.py::test_receive_works_sync - pluggy.PluggyTeardownRaisedWarning: A plugin raised an exception during an old-style hookwrapper teardown.
Plugin: unraisableexception, Hook: pytest_runtest_setup
PytestUnraisableExceptionWarning: Exception ignored in: <function _ConnectionRecord.checkout.<locals>.<lambda> at 0x7fd7dbd8b7e0>

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.1/x64/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 732, in <lambda>
    lambda ref: _finalize_fairy(
                ^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.1/x64/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 1034, in _finalize_fairy
    util.warn(message)
  File "/opt/hostedtoolcache/Python/3.12.1/x64/lib/python3.12/site-packages/sqlalchemy/util/langhelpers.py", line 1822, in warn
    _warnings_warn(msg, exc.SAWarning)
  File "/opt/hostedtoolcache/Python/3.12.1/x64/lib/python3.12/site-packages/sqlalchemy/util/langhelpers.py", line 1897, in _warnings_warn
    warnings.warn(message, category, stacklevel=stacklevel + 1)
sqlalchemy.exc.SAWarning: The garbage collector is trying to clean up non-checked-in connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7fd8245a8f50>>, which will be terminated.  Please ensure that SQLAlchemy pooled connections are returned to the pool explicitly, either by calling ``close()`` or by using appropriate context managers to manage their lifecycle.

For more information see https://pluggy.readthedocs.io/en/stable/api_reference.html#pluggy.PluggyTeardownRaisedWarning
```

Resolving patch:
diff --git a/src/prefect/server/database/configurations.py b/src/prefect/server/database/configurations.py
--- a/src/prefect/server/database/configurations.py
+++ b/src/prefect/server/database/configurations.py
@@ -23,6 +23,8 @@
     "SQLITE_BEGIN_MODE", default=None
 )
 
+ENGINES: Dict[Tuple[AbstractEventLoop, str, bool, float], AsyncEngine] = {}
+
 
 class BaseDatabaseConfiguration(ABC):
     """
@@ -91,8 +93,6 @@ async def begin_transaction(
 
 
 class AsyncPostgresConfiguration(BaseDatabaseConfiguration):
-    ENGINES: Dict[Tuple[AbstractEventLoop, str, bool, float], AsyncEngine] = {}
-
     async def engine(self) -> AsyncEngine:
         """Retrieves an async SQLAlchemy engine.
 
@@ -116,7 +116,7 @@ async def engine(self) -> AsyncEngine:
             self.echo,
             self.timeout,
         )
-        if cache_key not in self.ENGINES:
+        if cache_key not in ENGINES:
             # apply database timeout
             kwargs = dict()
             connect_args = dict()
@@ -151,9 +151,9 @@ async def engine(self) -> AsyncEngine:
                 **kwargs,
             )
 
-            self.ENGINES[cache_key] = engine
+            ENGINES[cache_key] = engine
             await self.schedule_engine_disposal(cache_key)
-        return self.ENGINES[cache_key]
+        return ENGINES[cache_key]
 
     async def schedule_engine_disposal(self, cache_key):
         """
@@ -173,7 +173,7 @@ async def schedule_engine_disposal(self, cache_key):
         """
 
         async def dispose_engine(cache_key):
-            engine = self.ENGINES.pop(cache_key, None)
+            engine = ENGINES.pop(cache_key, None)
             if engine:
                 await engine.dispose()
 
@@ -214,7 +214,6 @@ def is_inmemory(self) -> Literal[False]:
 
 
 class AioSqliteConfiguration(BaseDatabaseConfiguration):
-    ENGINES: Dict[Tuple[AbstractEventLoop, str, bool, float], AsyncEngine] = {}
     MIN_SQLITE_VERSION = (3, 24, 0)
 
     async def engine(self) -> AsyncEngine:
@@ -249,7 +248,7 @@ async def engine(self) -> AsyncEngine:
             self.echo,
             self.timeout,
         )
-        if cache_key not in self.ENGINES:
+        if cache_key not in ENGINES:
             # apply database timeout
             if self.timeout is not None:
                 kwargs["connect_args"] = dict(timeout=self.timeout)
@@ -269,9 +268,9 @@ async def engine(self) -> AsyncEngine:
             sa.event.listen(engine.sync_engine, "connect", self.setup_sqlite)
             sa.event.listen(engine.sync_engine, "begin", self.begin_sqlite_stmt)
 
-            self.ENGINES[cache_key] = engine
+            ENGINES[cache_key] = engine
             await self.schedule_engine_disposal(cache_key)
-        return self.ENGINES[cache_key]
+        return ENGINES[cache_key]
 
     async def schedule_engine_disposal(self, cache_key):
         """
@@ -291,7 +290,7 @@ async def schedule_engine_disposal(self, cache_key):
         """
 
         async def dispose_engine(cache_key):
-            engine = self.ENGINES.pop(cache_key, None)
+            engine = ENGINES.pop(cache_key, None)
             if engine:
                 await engine.dispose()
 
diff --git a/src/prefect/server/database/interface.py b/src/prefect/server/database/interface.py
--- a/src/prefect/server/database/interface.py
+++ b/src/prefect/server/database/interface.py
@@ -2,6 +2,7 @@
 from contextlib import asynccontextmanager
 
 import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 from prefect.server.database.alembic_commands import alembic_downgrade, alembic_upgrade
 from prefect.server.database.configurations import BaseDatabaseConfiguration
@@ -78,7 +79,7 @@ async def is_db_connectable(self):
         except Exception:
             return False
 
-    async def engine(self):
+    async def engine(self) -> AsyncEngine:
         """
         Provides a SqlAlchemy engine against a specific database.
         """


Test patch:
diff --git a/tests/fixtures/database.py b/tests/fixtures/database.py
--- a/tests/fixtures/database.py
+++ b/tests/fixtures/database.py
@@ -11,7 +11,11 @@
 from prefect.filesystems import LocalFileSystem
 from prefect.infrastructure import DockerContainer, Process
 from prefect.server import models, schemas
-from prefect.server.database.dependencies import provide_database_interface
+from prefect.server.database.configurations import ENGINES
+from prefect.server.database.dependencies import (
+    PrefectDBInterface,
+    provide_database_interface,
+)
 from prefect.server.orchestration.rules import (
     FlowOrchestrationContext,
     TaskOrchestrationContext,
@@ -27,12 +31,20 @@ def db(test_database_connection_url, safety_check_settings):
 
 
 @pytest.fixture(scope="session", autouse=True)
-async def database_engine(db):
+async def database_engine(db: PrefectDBInterface):
     """Produce a database engine"""
     engine = await db.engine()
-    try:
-        yield engine
-    finally:
+
+    yield engine
+
+    # At the end of the test session, dispose of _any_ open engines, not just the one
+    # that we produced here.  Other engines may have been created from different event
+    # loops, and we need to ensure that they are all disposed of.
+
+    engines = list(ENGINES.values())
+    ENGINES.clear()
+
+    for engine in engines:
         await engine.dispose()
 
 
diff --git a/tests/utilities/test_asyncutils.py b/tests/utilities/test_asyncutils.py
--- a/tests/utilities/test_asyncutils.py
+++ b/tests/utilities/test_asyncutils.py
@@ -381,10 +381,13 @@ async def run_test():
         await add_event_loop_shutdown_callback(set_event)
 
     loop = asyncio.new_event_loop()
-    thread = threading.Thread(target=loop.run_until_complete(run_test()))
-    thread.start()
-    assert not callback_called.wait(timeout=1)
-    thread.join(timeout=1)
+    try:
+        thread = threading.Thread(target=loop.run_until_complete(run_test()))
+        thread.start()
+        assert not callback_called.wait(timeout=1)
+        thread.join(timeout=1)
+    finally:
+        loop.close()
 
 
 async def test_gather():