Fix connection pool lifecycle to stop leaking connections

src/postgres_mcp/server.py

@@ -554,6 +554,30 @@ async def get_top_queries(
         return format_error_response(str(e))
 
 
+async def exit_if_orphaned(poll_interval: float = 2.0) -> None:
+    """Force-exit once this process has been orphaned (reparented to init, PID 1).
+
+    Used only for the stdio transport, where the server is a child of the MCP
+    client. If the client exits without closing stdin (e.g. it was force-killed),
+    the OS reparents us to PID 1 and run_stdio_async() never returns, so the
+    server would linger forever holding connections against the database role.
+
+    We cannot tear this down by cancelling run_stdio_async() — its stdin reader
+    blocks in a thread and does not unwind on cancellation. So once orphaned we
+    release the pool and hard-exit the process. The pool is closed first, so the
+    connections are returned to the server cleanly before we go.
+    """
+    while os.getppid() != 1:
+        await asyncio.sleep(poll_interval)
+    logger.info("Parent process exited; shutting down orphaned server")
+    try:
+        await db_connection.close()
+        logger.info("Closed database connections")
+    except Exception as e:
+        logger.error(f"Error closing database connections: {e}")
+    os._exit(0)
+
+
 async def main():
     # Parse command line arguments
     parser = argparse.ArgumentParser(description="PostgreSQL MCP Server")
@@ -656,17 +680,41 @@ async def main():
         logger.warning("Signal handling not supported on Windows")
         pass
 
-    # Run the server with the selected transport (always async)
-    if args.transport == "stdio":
-        await mcp.run_stdio_async()
-    elif args.transport == "sse":
-        mcp.settings.host = args.sse_host
-        mcp.settings.port = args.sse_port
-        await mcp.run_sse_async()
-    elif args.transport == "streamable-http":
-        mcp.settings.host = args.streamable_http_host
-        mcp.settings.port = args.streamable_http_port
-        await mcp.run_streamable_http_async()
+    # Run the server with the selected transport (always async).
+    #
+    # An MCP client normally signals shutdown by closing stdin,
+    # which makes run_stdio_async() return and the finally below releases the
+    # pool. But if the client is force-killed, stdin EOF may never arrive and the
+    # process is reparented to init (PID 1), lingering forever and pinning its
+    # connections. exit_if_orphaned() runs alongside the transport and hard-exits
+    # in that case (run_stdio_async cannot be cancelled cleanly).
+    orphan_task = None
+    try:
+        if args.transport == "stdio":
+            orphan_task = asyncio.ensure_future(exit_if_orphaned())
+            await mcp.run_stdio_async()
+        elif args.transport == "sse":
+            mcp.settings.host = args.sse_host
+            mcp.settings.port = args.sse_port
+            await mcp.run_sse_async()
+        elif args.transport == "streamable-http":
+            mcp.settings.host = args.streamable_http_host
+            mcp.settings.port = args.streamable_http_port
+            await mcp.run_streamable_http_async()
+    finally:
+        # Clean-exit path (stdin EOF / signal): stop the watchdog and release the pool.
+        # (On the orphan path the watchdog calls os._exit and we never reach here.)
+        if orphan_task is not None:
+            orphan_task.cancel()
+            try:
+                await orphan_task
+            except asyncio.CancelledError:
+                pass
+        try:
+            await db_connection.close()
+            logger.info("Closed database connections")
+        except Exception as e:
+            logger.error(f"Error closing database connections: {e}")
 
 
 async def shutdown(sig=None):

src/postgres_mcp/sql/sql_driver.py

@@ -85,11 +85,18 @@ async def pool_connect(self, connection_url: Optional[str] = None) -> AsyncConne
         await self.close()
 
         try:
-            # Configure connection pool with appropriate settings
+            # Configure connection pool with appropriate settings.
+            #
+            # min_size=0: an idle server holds ZERO connections. With min_size>=1
+            # every running (or orphaned) server permanently pins a connection
+            # against the database role, so a handful of leftover processes can
+            # exhaust a low per-role connection limit. max_idle reaps connections
+            # that go unused, so the pool drifts back to zero between queries.
             self.pool = AsyncConnectionPool(
                 conninfo=url,
-                min_size=1,
-                max_size=5,
+                min_size=0,
+                max_size=3,
+                max_idle=60,  # seconds; release idle connections back to the server
                 open=False,  # Don't connect immediately, let's do it explicitly
             )
 

tests/unit/test_lifecycle.py

@@ -0,0 +1,113 @@
+"""Tests for the stdio connection-pool lifecycle: the orphan watchdog and the
+guarantee that the pool is released when the server exits.
+
+These cover the behavior added to stop the server leaking connections against
+the database role when its MCP client is force-killed without closing stdin —
+see exit_if_orphaned() and main() in postgres_mcp.server.
+"""
+
+import asyncio
+import sys
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_exit_if_orphaned_closes_pool_before_exiting():
+    """Once orphaned (getppid()==1) the watchdog must close the pool and THEN
+    hard-exit — closing first is what returns the connections cleanly."""
+    from postgres_mcp import server
+
+    order = []
+    close_mock = AsyncMock(side_effect=lambda: order.append("close"))
+    exit_mock = MagicMock(side_effect=lambda code: order.append(("exit", code)))
+
+    with (
+        patch("postgres_mcp.server.os.getppid", return_value=1),
+        patch("postgres_mcp.server.db_connection.close", close_mock),
+        patch("postgres_mcp.server.os._exit", exit_mock),
+    ):
+        await server.exit_if_orphaned(poll_interval=0.001)
+
+    assert order == ["close", ("exit", 0)]
+
+
+@pytest.mark.asyncio
+async def test_exit_if_orphaned_waits_until_orphaned():
+    """While the parent is alive (getppid()!=1) the watchdog keeps polling and
+    must NOT close the pool or exit."""
+    from postgres_mcp import server
+
+    ppids = iter([1234, 1234, 1])  # alive twice, then reparented to init
+
+    def fake_getppid():
+        return next(ppids)
+
+    close_mock = AsyncMock()
+    exit_mock = MagicMock()
+
+    with (
+        patch("postgres_mcp.server.os.getppid", side_effect=fake_getppid),
+        patch("postgres_mcp.server.db_connection.close", close_mock),
+        patch("postgres_mcp.server.os._exit", exit_mock),
+    ):
+        await server.exit_if_orphaned(poll_interval=0.001)
+
+    close_mock.assert_awaited_once()
+    exit_mock.assert_called_once_with(0)
+
+
+async def _run_main_capturing_tasks(transport, created):
+    """Run main() for the given transport, recording every task spawned via
+    asyncio.ensure_future so tests can assert whether the watchdog was started."""
+    from postgres_mcp.server import main
+
+    real_ensure_future = asyncio.ensure_future
+
+    def spy(coro, *args, **kwargs):
+        task = real_ensure_future(coro, *args, **kwargs)
+        created.append(task)
+        return task
+
+    original_argv = sys.argv
+    sys.argv = ["postgres_mcp", "postgresql://user:password@localhost/db", f"--transport={transport}"]
+    try:
+        with (
+            # never actually orphaned during the test
+            patch("postgres_mcp.server.os.getppid", return_value=99999),
+            patch("postgres_mcp.server.db_connection.pool_connect", AsyncMock()),
+            patch("postgres_mcp.server.db_connection.close", AsyncMock()),
+            patch("postgres_mcp.server.mcp.run_stdio_async", AsyncMock()),
+            patch("postgres_mcp.server.mcp.run_sse_async", AsyncMock()),
+            patch("postgres_mcp.server.mcp.run_streamable_http_async", AsyncMock()),
+            patch("postgres_mcp.server.asyncio.ensure_future", side_effect=spy),
+        ):
+            await main()
+    finally:
+        sys.argv = original_argv
+
+
+@pytest.mark.asyncio
+async def test_stdio_starts_watchdog_and_cancels_it_on_clean_exit():
+    """stdio must start exactly one watchdog task and cancel it when the
+    transport returns cleanly (stdin EOF), leaving no task running."""
+    created = []
+    await _run_main_capturing_tasks("stdio", created)
+
+    assert len(created) == 1, "stdio should start exactly one watchdog task"
+    with pytest.raises(asyncio.CancelledError):
+        await created[0]  # deterministically confirms it was cancelled
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("transport", ["sse", "streamable-http"])
+async def test_non_stdio_transports_do_not_start_watchdog(transport):
+    """The orphan watchdog is stdio-only; long-running HTTP transports manage
+    their own lifecycle and must not spawn it."""
+    created = []
+    await _run_main_capturing_tasks(transport, created)
+
+    assert created == [], f"{transport} must not start the orphan watchdog"

tmp/fix-connection-pool-lifecycle-2026-06-05-review.md

@@ -0,0 +1,28 @@
+# Code Review Exchange — postgres-mcp connection-pool lifecycle fix
+
+Independent reviewer: auggie / gpt-5.5 (whole-codebase index), 3 passes.
+
+## Consolidated Review (from 3 independent passes)
+
+- Pass 1: `ready` — LGTM, no issues.
+- Pass 2: `ready` — LGTM, no issues.
+- Pass 3: `needs work` — one finding (below).
+
+### Finding 1
+- **File**: src/postgres_mcp/server.py (557-578, run block + _exit_if_orphaned), src/postgres_mcp/sql/sql_driver.py (pool)
+- **Lens**: completeness & coverage
+- **Severity**: medium
+- **Confidence**: Flagged in 1/3 passes (the other two passed it clean)
+- **Issue**: The riskiest new behavior — a stdio-only orphan watchdog that closes the global pool then calls os._exit(0) — had no checked-in test. Existing transport tests only assert which run_*_async() is called; nothing proved the watchdog exits after getppid()==1, closes the pool before exiting, is cancelled on clean stdio return, and is never started for sse/streamable-http. Only manual/empirical verification existed.
+- **Suggestion**: Add focused unit tests patching os.getppid/os._exit/db_connection.close to exercise _exit_if_orphaned (asserting close-before-exit), plus main() tests asserting the watchdog is created+cancelled for stdio and not created for sse/streamable-http, without invoking real os._exit.
+
+## Response by Author
+
+### Re: Finding 1
+- **Status**: fixed
+- **Response**: Valid — this is exactly the non-obvious lifecycle logic that warrants real coverage, especially heading into an upstream PR. Added tests/unit/test_lifecycle.py (5 tests): close-before-exit ordering, poll-until-orphaned (no premature close/exit), stdio starts exactly one watchdog and cancels it on clean exit, and sse/streamable-http start none. Also tightened the finally to await the cancelled watchdog (no dangling task) — confirmed it doesn't break the 6 existing transport tests.
+- **Changes**: src/postgres_mcp/server.py (await watchdog cancellation in finally); tests/unit/test_lifecycle.py (new). `uv run pytest tests/unit/test_lifecycle.py tests/unit/test_transport.py` → 11 passed.
+
+No findings disputed. The two LGTM passes plus the addressed coverage gap → resolved.
+
+<!-- AGREED -->