[mangeshraut712]

Fixes #4826

Summary

Changes

When PG-wire startup validation ran under asyncio.run(), the pool and aiorwlock were bound to a temporary loop that then closed. mcp.run() serves tools on a different loop, so the first run_query failed with a different-event-loop RuntimeError.

Move PG-wire (PsycopgPoolConnection) startup validation into the MCP server lifespan so the pool initializes on the serving loop, and clean up loop-local connection maps on shutdown. Non-PG-wire validation stays on the existing sync path.

User experience

Before: the first run_query after PG-wire startup could fail with a closed / different event loop error. After: validation runs on the serving loop, so the first query uses a live pool.

Checklist

• I have reviewed the contributing guidelines
• I have performed a self-review of this change
• Changes have been tested
• Changes are documented

Tests: regression covering main() → lifespan validation → first run_query with a fake wire pool. uv run pytest — 2501 passed in the server package context. Ruff and pyright clean on touched files.

Is this a breaking change? (Y/N) N

Acknowledgment

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.

[mangeshraut712]

Fixes #4826

Why: PG-wire startup validation under asyncio.run() bound the pool to a temporary loop that then closed; first run_query failed on the serving loop.

What: Move PG-wire (PsycopgPoolConnection) startup validation into MCP server lifespan so the pool initializes on the serving loop; clean up loop-local maps on shutdown.

Test: Regression main() → lifespan validation → first run_query with fake wire pool. Package pytest green; ruff/pyright clean on touched files.

[mangeshraut712]

CI is green. Avoids binding the PG-wire connection to a closed event loop. Please review and merge when you can.

[codecov]

Codecov Report

❌ Patch coverage is 78.04878% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.68%. Comparing base (465295e) to head (4c142d8).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...s-mcp-server/awslabs/postgres_mcp_server/server.py	78.12%	6 Missing and 1 partial ⚠️
...ostgres_mcp_server/connection/db_connection_map.py	77.77%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4846   +/-   ##
=======================================
  Coverage   93.68%   93.68%           
=======================================
  Files        1078     1078           
  Lines       93123    93154   +31     
  Branches    15087    15090    +3     
=======================================
+ Hits        87238    87275   +37     
+ Misses       3505     3494   -11     
- Partials     2380     2385    +5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[kennthhz]

Diagnosis looks right to me, and the lifespan is the correct mechanism. I traced the root cause and it holds: every pool operation in PsycopgPoolConnection goes through self.rw_lock, and aiorwlock binds its loop lazily on first acquire, then raises on cross-loop use. asyncio.run() bound it to a loop that then closed.

Two things I'd want addressed before this merges, plus some smaller notes inline.

The one I care most about: nothing appears to test that the lifespan is actually registered. The new test calls server_lifespan(...) directly and patches mcp.run, so lifespan=server_lifespan in the constructor is unasserted. That matters more than it looks — see the inline comment.

I verified the framework behavior against mcp 2.0.0 and aiorwlock 1.5.0 locally. One claim I could not check: whether a test elsewhere in the suite already covers the lifespan registration. If so, ignore that comment.

[kennthhz]

This registration is what makes the whole fix work, and I don't think anything asserts it.

The new test enters server_lifespan(...) directly and patches mcp.run wholesale, so deleting this one kwarg leaves the suite green. Consider what the server then does: main() sets startup_pg_wire_connection, the lifespan never runs, and the first run_query initializes the pool lazily on the serving loop — which works, because check_expiry() falls through to initialize_pool() when self.pool is None.

So the reported bug stays fixed, and what silently disappears instead is startup validation — including the --privilege_check strict guardrail that refuses to start against an over-privileged connection.

Worth noting what that implies about the change: the minimal fix for #4826 is just "stop validating on a throwaway loop." The lifespan is here to preserve validation while fixing it. So the property to pin isn't "the first query succeeds," it's "validation ran, on the serving loop, before serving." An assertion that mcp is constructed with this lifespan, or one test driven through the real mcp.run(), would cover it.

[kennthhz]

This changes the failure mode for a deliberate refusal. Previously ConnectionValidationError meant logger.error('Refusing to start: ...') then sys.exit(1) — one clean line. Raising from inside the lifespan sends it through anyio's task group, and anyio 4.x wraps even a single exception:

type: ExceptionGroup
is Boom? False
wrapped: ['Boom']

So an operator who set --privilege_check strict and connected as a superuser now gets an ExceptionGroup traceback for an expected policy decision. The Data API path above still exits cleanly for the same condition.

Note this also rules out the obvious fix: a plain except ConnectionValidationError around mcp.run() won't catch it, since isinstance is False, and this package supports Python 3.10 where except* isn't available. Simplest option is to handle it here and sys.exit(1) before it ever reaches the task group.

[kennthhz]

Minor: consider yield {} with AsyncIterator[Dict[str, Any]].

The lowlevel default lifespan yields {}, and lifespan_wrapper passes the yielded value straight through as lifespan_state. So this changes ctx.request_context.lifespan_context from an empty dict to None, and anything doing .get(...) on it would hit an AttributeError.

Nothing reads lifespan_context in this server today, so this is a latent trap rather than a bug — just cheap to close.

[kennthhz]

This isinstance check is now load-bearing for lifecycle policy, which forks startup into two paths with three consequences: the failure UX diverges by connection type (clean sys.exit(1) here vs. an ExceptionGroup traceback via the lifespan), the ConnectionValidationError/Exception handler is duplicated in two places and will drift, and a future connection type with loop-affine state silently takes the old broken path.

Validating every connection in the lifespan and deleting this branch would collapse all three. If the sync path has to stay, putting the policy on the connection — something like AbstractDBConnection.requires_serving_loop — would at least keep it from being an isinstance at the call site.

[kennthhz]

Reads correctly, and taking the snapshot under the lock before awaiting is the right shape.

One thing worth tightening in the sibling close_all() while this is fresh: its except RuntimeError branch assumes the only possible RuntimeError is "already inside a running loop," and discards the coroutine with no log. It also catches the cross-loop binding error this PR is about. I tested that path against a lock bound to a dead loop — the pool does get closed, because the error surfaces on lock release after the close already took effect, and then gets swallowed silently. Harmless today, purely by luck. If aiorwlock ever raised on acquire instead, the close would be skipped with zero diagnostics.

[kennthhz]

The map is cleared before any close is attempted, so a connection whose close() raises is already gone from the map — no retry, nothing to inspect, just the warning below. Mostly fine at shutdown, but dropping it from the map only after a successful close would make a stuck pool visible.

[kennthhz]

Closes run serially with no timeout. pool.close() can block waiting on checked-out connections, so one wedged pool blocks shutdown indefinitely. An asyncio.gather with a bounded wait would be cheap insurance.

Also: the warning below interpolates key, which carries the endpoint hostname into the logs. Consistent with existing logging here, so flagging rather than asking for a change.

[kennthhz]

Calling server_lifespan directly is what leaves the constructor wiring untested — see my comment on the lifespan=server_lifespan line.

One suggestion for the test itself: its fake Pool is entirely loop-agnostic, so all of the protection here comes from the real aiorwlock, not from the pool. It does work — reverting the isinstance branch binds the lock on one loop and raises on the other, failing this test. But a future refactor that drops the rw_lock would quietly remove the test's teeth with nothing to signal it. Worth a comment saying the lock is the thing under test.

[kennthhz]

AsyncMock() would be the idiomatic form here and would also let you assert it was awaited, rather than driving a MagicMock with an async side_effect.