diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..ae29b60b 100644 --- a/src/postgres_mcp/server.py +++ b/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): diff --git a/src/postgres_mcp/sql/sql_driver.py b/src/postgres_mcp/sql/sql_driver.py index 5beacb03..e6d095fa 100644 --- a/src/postgres_mcp/sql/sql_driver.py +++ b/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 ) diff --git a/tests/unit/test_lifecycle.py b/tests/unit/test_lifecycle.py new file mode 100644 index 00000000..84473165 --- /dev/null +++ b/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" diff --git a/tmp/fix-connection-pool-lifecycle-2026-06-05-review.md b/tmp/fix-connection-pool-lifecycle-2026-06-05-review.md new file mode 100644 index 00000000..3b0b54c8 --- /dev/null +++ b/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. + +