mirror of
https://github.com/apache/superset.git
synced 2026-05-29 20:29:34 +00:00
fix(mcp): handle SSL connection drop during pre-call session teardown
When RDS drops an SSL connection due to idle timeout or max-connection-age, `db.session.remove()` in `sync_wrapper` raises `OperationalError` because the implicit rollback inside `session.close()` fails on the dead DBAPI connection. This caused the MCP tool call to fail even when the operation itself completed successfully, and left a dead connection in the pool. Introduce `_remove_session_safe()` which: - Catches `OperationalError` from `session.remove()` (SSL/network errors) - Calls `session.invalidate()` to mark the dead connection for pool discard - Retries `session.remove()` so the scoped registry is clean before the tool runs Replace the bare `db.session.remove()` in `sync_wrapper` with `_remove_session_safe()`. Add a unit test verifying `invalidate()` is called and remove is retried on SSL error.
This commit is contained in:
@@ -352,9 +352,9 @@ def test_mcp_auth_hook_preserves_g_user_in_request_context(app) -> None:
|
||||
|
||||
def _assert_preserved_then_return():
|
||||
"""Verify g.user was preserved (not cleared) before returning."""
|
||||
assert hasattr(g, "user"), (
|
||||
"g.user should be preserved in request context but was removed"
|
||||
)
|
||||
assert hasattr(
|
||||
g, "user"
|
||||
), "g.user should be preserved in request context but was removed"
|
||||
assert g.user is middleware_user, (
|
||||
"g.user should be preserved in request context but was changed; "
|
||||
f"g.user={g.user}"
|
||||
@@ -409,6 +409,53 @@ def test_mcp_auth_hook_removes_stale_db_session_in_sync_wrapper(app) -> None:
|
||||
assert result == "fresh"
|
||||
|
||||
|
||||
def test_sync_wrapper_handles_ssl_error_on_pre_call_remove(app) -> None:
|
||||
"""sync_wrapper tolerates OperationalError from db.session.remove() before the call.
|
||||
|
||||
If the underlying DBAPI connection died between requests (e.g. RDS SSL
|
||||
idle-timeout), the rollback implicit in session.close() raises
|
||||
OperationalError. _remove_session_safe() should:
|
||||
- Log a warning
|
||||
- Call session.invalidate() to mark the dead connection for pool discard
|
||||
- Retry session.remove() so the registry is clean
|
||||
- Allow the tool to run successfully
|
||||
"""
|
||||
from sqlalchemy.exc import OperationalError as SAOperationalError
|
||||
|
||||
fresh_user = _make_mock_user("fresh")
|
||||
|
||||
def dummy_tool() -> str:
|
||||
"""Dummy sync tool."""
|
||||
return g.user.username
|
||||
|
||||
wrapped = mcp_auth_hook(dummy_tool)
|
||||
|
||||
remove_call_count = 0
|
||||
|
||||
def _flaky_remove() -> None:
|
||||
nonlocal remove_call_count
|
||||
remove_call_count += 1
|
||||
if remove_call_count == 1:
|
||||
raise SAOperationalError(
|
||||
"SSL connection has been closed unexpectedly", None, None
|
||||
)
|
||||
|
||||
with app.test_request_context():
|
||||
g.user = fresh_user
|
||||
with patch("superset.extensions.db") as mock_db:
|
||||
mock_db.session.remove.side_effect = _flaky_remove
|
||||
|
||||
with patch(
|
||||
"superset.mcp_service.auth.get_user_from_request",
|
||||
return_value=fresh_user,
|
||||
):
|
||||
result = wrapped()
|
||||
|
||||
assert result == "fresh"
|
||||
assert mock_db.session.invalidate.called, "invalidate() must be called on SSL error"
|
||||
assert remove_call_count == 2, "remove() must be retried after SSL error"
|
||||
|
||||
|
||||
# -- default_user_resolver --
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user