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
```
Even after closing out remaining engines at the end of the test suite, there is still at least one open connection to the DB as the test session exits.

https://github.com/PrefectHQ/prefect/actions/runs/7829266132/attempts/1

```
ERROR tests/runner/test_webserver.py::TestWebserverSettings::test_webserver_settings_are_respected - pytest.PytestUnraisableExceptionWarning: Exception ignored in: <socket.socket fd=-1, family=AddressFamily.AF_INET6, type=SocketKind.SOCK_STREAM, proto=6>

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/abc.py", line 102, in __subclasscheck__
    return _abc_subclasscheck(cls, subclass)
ResourceWarning: unclosed <socket.socket fd=27, family=AddressFamily.AF_INET6, type=SocketKind.SOCK_STREAM, proto=6, laddr=('::1', 49296, 0, 0), raddr=('::1', 5432, 0, 0)>
```
I think the new approach will require instrumenting the connection pool to track all checkouts and their corresponding checkins, so that we can close any remaining connections.  As a start, attach to the relevant checkin/out events and capture the connection and also the traceback of where it was checked out.  This may help find the ultimate source of this bug, but otherwise we can use this tracking to close all remaining open connections.

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
@@ -1,4 +1,5 @@
 import sqlite3
+import traceback
 from abc import ABC, abstractmethod
 from asyncio import AbstractEventLoop, get_running_loop
 from contextlib import asynccontextmanager
@@ -26,6 +27,71 @@
 ENGINES: Dict[Tuple[AbstractEventLoop, str, bool, float], AsyncEngine] = {}
 
 
+class ConnectionTracker:
+    """A test utility which tracks the connections given out by a connection pool, to
+    make it easy to see which connections are currently checked out and open."""
+
+    all_connections: Dict[sa.AdaptedConnection, str]
+    open_connections: Dict[sa.AdaptedConnection, str]
+    left_field_closes: Dict[sa.AdaptedConnection, str]
+    connects: int
+    closes: int
+    active: bool
+
+    def __init__(self) -> None:
+        self.active = False
+        self.all_connections = {}
+        self.open_connections = {}
+        self.left_field_closes = {}
+        self.connects = 0
+        self.closes = 0
+
+    def track_pool(self, pool: sa.pool.Pool):
+        sa.event.listen(pool, "connect", self.on_connect)
+        sa.event.listen(pool, "close", self.on_close)
+        sa.event.listen(pool, "close_detached", self.on_close_detached)
+
+    def on_connect(
+        self,
+        adapted_connection: sa.AdaptedConnection,
+        connection_record: sa.pool.ConnectionPoolEntry,
+    ):
+        self.all_connections[adapted_connection] = traceback.format_stack()
+        self.open_connections[adapted_connection] = traceback.format_stack()
+        self.connects += 1
+
+    def on_close(
+        self,
+        adapted_connection: sa.AdaptedConnection,
+        connection_record: sa.pool.ConnectionPoolEntry,
+    ):
+        try:
+            del self.open_connections[adapted_connection]
+        except KeyError:
+            self.left_field_closes[adapted_connection] = traceback.format_stack()
+        self.closes += 1
+
+    def on_close_detached(
+        self,
+        adapted_connection: sa.AdaptedConnection,
+    ):
+        try:
+            del self.open_connections[adapted_connection]
+        except KeyError:
+            self.left_field_closes[adapted_connection] = traceback.format_stack()
+        self.closes += 1
+
+    def clear(self):
+        self.all_connections.clear()
+        self.open_connections.clear()
+        self.left_field_closes.clear()
+        self.connects = 0
+        self.closes = 0
+
+
+TRACKER: ConnectionTracker = ConnectionTracker()
+
+
 class BaseDatabaseConfiguration(ABC):
     """
     Abstract base class used to inject database connection configuration into Prefect.
@@ -151,6 +217,9 @@ async def engine(self) -> AsyncEngine:
                 **kwargs,
             )
 
+            if TRACKER.active:
+                TRACKER.track_pool(engine.pool)
+
             ENGINES[cache_key] = engine
             await self.schedule_engine_disposal(cache_key)
         return ENGINES[cache_key]
@@ -268,6 +337,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)
 
+            if TRACKER.active:
+                TRACKER.track_pool(engine.pool)
+
             ENGINES[cache_key] = engine
             await self.schedule_engine_disposal(cache_key)
         return ENGINES[cache_key]


Test patch:
diff --git a/tests/cli/deployment/test_deployment_cli.py b/tests/cli/deployment/test_deployment_cli.py
--- a/tests/cli/deployment/test_deployment_cli.py
+++ b/tests/cli/deployment/test_deployment_cli.py
@@ -264,7 +264,7 @@ def test_set_schedule_interval_without_anchor_date(self, flojo):
                 "inspect",
                 "rence-griffith/test-deployment",
             ],
-            expected_output_contains=["10.76"],  # 100 m record
+            expected_output_contains=["'interval': 10.76,"],  # 100 m record
             expected_code=0,
         )
 
@@ -285,7 +285,9 @@ def test_set_schedule_interval_without_anchor_date(self, flojo):
                 "inspect",
                 "rence-griffith/test-deployment",
             ],
-            expected_output_contains=["10.49"],  # flo-jo breaks the world record
+            expected_output_contains=[
+                "'interval': 10.49,"
+            ],  # flo-jo breaks the world record
             expected_code=0,
         )
 
@@ -296,7 +298,7 @@ def test_set_no_schedule(self, flojo):
                 "inspect",
                 "rence-griffith/test-deployment",
             ],
-            expected_output_contains=["10.76"],  # 100 m record
+            expected_output_contains=["'interval': 10.76,"],  # 100 m record
             expected_code=0,
         )
 
@@ -317,7 +319,7 @@ def test_set_no_schedule(self, flojo):
                 "inspect",
                 "rence-griffith/test-deployment",
             ],
-            expected_output_does_not_contain=["10.76"],
+            expected_output_does_not_contain=["'interval': 10.76,"],
             expected_output_contains=["'schedule': None"],
             expected_code=0,
         )
diff --git a/tests/fixtures/database.py b/tests/fixtures/database.py
--- a/tests/fixtures/database.py
+++ b/tests/fixtures/database.py
@@ -1,7 +1,10 @@
 import asyncio
 import datetime
+import gc
 import warnings
 
+import aiosqlite
+import asyncpg
 import pendulum
 import pytest
 from sqlalchemy.exc import InterfaceError
@@ -11,7 +14,7 @@
 from prefect.filesystems import LocalFileSystem
 from prefect.infrastructure import DockerContainer, Process
 from prefect.server import models, schemas
-from prefect.server.database.configurations import ENGINES
+from prefect.server.database.configurations import ENGINES, TRACKER
 from prefect.server.database.dependencies import (
     PrefectDBInterface,
     provide_database_interface,
@@ -33,6 +36,8 @@ def db(test_database_connection_url, safety_check_settings):
 @pytest.fixture(scope="session", autouse=True)
 async def database_engine(db: PrefectDBInterface):
     """Produce a database engine"""
+    TRACKER.active = True
+
     engine = await db.engine()
 
     yield engine
@@ -47,6 +52,22 @@ async def database_engine(db: PrefectDBInterface):
     for engine in engines:
         await engine.dispose()
 
+    # Now confirm that after disposing all engines, all connections are closed
+
+    for connection in TRACKER.all_connections:
+        driver_connection = connection.driver_connection
+        if isinstance(driver_connection, asyncpg.Connection):
+            assert driver_connection.is_closed()
+        elif isinstance(driver_connection, aiosqlite.Connection):
+            assert not driver_connection._connection
+
+    # Finally, free up all references to connections and clean up proactively so that
+    # we don't have any lingering connections after this.  This should prevent
+    # post-test-session ResourceWarning errors about unclosed connections.
+
+    TRACKER.clear()
+    gc.collect()
+
 
 @pytest.fixture
 def print_query(database_engine):