This is an automated email from the ASF dual-hosted git repository. rusackas pushed a commit to branch fix/mcp-worker-metadata-sizing in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6d612e9768b066ef49b035bc8c9c6c0350b5db94 Author: Evan Rusackas <[email protected]> AuthorDate: Wed Sep 30 18:04:33 2026 -0700 fix(mcp): bound metadata-only workers by remaining pool capacity, plus a review pass tool_worker_count's own math reserves exactly one spare metadata connection beyond what warehouse workers and their cancellations need (2 * workers + 1). metadata_tool_worker_count was sized completely independently of that budget (up to 16 by default, unconditionally), so under load the two bounds combined could check out more metadata connections than the pool actually has -- reopening the exact exhaustion this PR exists to fix. metadata_tool_worker_count is now capped by whatever's left of the pool after the warehouse reservation, falling back to sharing the warehouse pool (metadata_size=0) rather than ever overdrawing it. Also closes the superset/sql/ 100% coverage gap: cancellable_engine's unconditional admit() call on a connect-event prequery (the only hook such a prequery can reach) had no test exercising it. Plus a review pass addressing what a deeper /code-review turned up: - on_list_tools and the two search-transform _get_visible_tools overrides called run_in_metadata_thread with no try/except at all. That helper reloads the caller's ORM user itself before the passed-in filter function ever runs, so a failure there (e.g. the metadata pool exhausted) propagated uncaught -- the filter's own fail-open handling never got a chance to run. All three now fail open the same way that handling already does; call-time RBAC still enforces permissions. - QueryCancellation.capture() never reset cancel_id to None before a re-capture attempt. A multi-statement call whose second capture failed kept register() republishing the first (already-finished) statement's handle instead of going unavailable for the second. - WorkerPool.cancel() only released its permit on RuntimeError; any other exception from submitting cancellation I/O leaked a permit for the life of the process. - run_in_worker's pool.submit() call was outside the try/finally that resets _mcp_user_id_var, so a busy-pool rejection left a stale user id from a previous call in that context for audit logging. - run_in_worker and TransportContext.forward both relabel any TimeoutError raised while awaiting as "our deadline fired," including one genuinely raised by the tool/driver itself (e.g. a socket-level read timeout, a TimeoutError subclass since Python 3.10). Both now check whether the awaited future is already done (and, for the unshielded forward() case, not cancelled) before relabeling it, so a real driver timeout propagates as itself instead of a misleading generic deadline message with the original exception discarded. Deliberately left alone, flagged rather than guessed at: - cancellation.py's before_cursor_execute listener can raise (WorkerDeadlineExceeded/ToolError via admission) from inside SQLAlchemy's own event dispatch, outside the try/except that routes genuine DBAPI errors through handle_error's cleanup. Fixing this safely needs a clearer picture of how SQLAlchemy's connection pool actually treats a non-DBAPI exception raised from this hook than I have confidence in -- flagging rather than guessing at a pool-hygiene fix. - METADATA_ONLY_TOOLS is a hand-maintained list disconnected from where tools declare their own metadata; restructuring that is more invasive than this PR's scope. - run_in_worker's execute() reloads g.user via a metadata round trip that JWT/API-key auth immediately overwrites; skipping it safely needs tracing all four auth priority paths in get_user_from_request, not just the common ones. All touched files pass pre-commit (mypy, ruff, pylint); mcp_service and sql/execution suites pass (6451 passed, the 3 pre-existing failures in test_time_range_validation.py and test_query_dataset.py::test_query_dataset_reexecutes_across_rollover reproduce identically without this commit); superset/sql/'s 100% coverage gate passes. Co-Authored-By: Evan Rusackas <[email protected]> Co-Authored-By: Claude Sonnet 5 <[email protected]> --- superset/mcp_service/middleware.py | 12 +- superset/mcp_service/server.py | 30 ++++- superset/mcp_service/worker.py | 138 ++++++++++++++++----- tests/unit_tests/mcp_service/test_middleware.py | 21 ++++ tests/unit_tests/mcp_service/test_worker.py | 54 +++++++- .../mcp_service/test_worker_metadata_pool.py | 51 +++++++- .../unit_tests/sql/execution/test_cancellation.py | 26 ++++ 7 files changed, 293 insertions(+), 39 deletions(-) diff --git a/superset/mcp_service/middleware.py b/superset/mcp_service/middleware.py index 4be6260f3fb..69c471eb025 100644 --- a/superset/mcp_service/middleware.py +++ b/superset/mcp_service/middleware.py @@ -1200,7 +1200,17 @@ class RBACToolVisibilityMiddleware(Middleware): # User and permission lookups need a metadata connection. Tool workers # can hold every pooled connection while they wait on this loop, so the # loop must never wait for one itself. - return await run_in_metadata_thread(self._visible_tools, tools) + try: + return await run_in_metadata_thread(self._visible_tools, tools) + except Exception: # noqa: BLE001 + # A failure setting up the metadata thread itself (e.g. the + # metadata pool is exhausted) happens before _visible_tools's own + # try/except ever runs. Fail open the same way that branch does; + # call-time RBAC still enforces permissions. + logger.warning( + "MCP tool list: failed to filter tools, showing all", exc_info=True + ) + return tools @staticmethod def _visible_tools(tools: list[Tool]) -> list[Tool]: diff --git a/superset/mcp_service/server.py b/superset/mcp_service/server.py index dc43ac5095c..51153ae3972 100644 --- a/superset/mcp_service/server.py +++ b/superset/mcp_service/server.py @@ -430,6 +430,28 @@ def _filter_tools_by_current_user_permission(tools: Sequence[Any]) -> list[Any]: return [tool for tool in tools if _tool_allowed_for_current_user(tool)] +async def _filter_visible_tools_fail_open(tools: Sequence[Any]) -> Sequence[Any]: + """Run the permission filter in the metadata thread, failing open on error. + + ``run_in_metadata_thread`` reloads the caller's ORM user itself before the + filter ever runs (e.g. a metadata-pool-exhaustion failure), so a bare + ``await run_in_metadata_thread(...)`` here would raise before any fail-open + handling inside the filter gets a chance to run. Call-time RBAC still + enforces permissions, so an unexpected failure here shows every tool + rather than breaking search, matching + ``RBACToolVisibilityMiddleware.on_list_tools``. + """ + try: + return await run_in_metadata_thread( + _filter_tools_by_current_user_permission, tools + ) + except Exception: # noqa: BLE001 + logger.warning( + "MCP tool search: failed to filter tools, showing all", exc_info=True + ) + return tools + + def _create_search_result_serializer( config: dict[str, Any], ) -> Any: @@ -671,9 +693,7 @@ def _create_search_transform( # noqa: C901 async def _get_visible_tools(self, ctx: Context) -> Sequence[Any]: """Return only tools visible to the current authenticated user.""" tools = await super()._get_visible_tools(ctx) - return await run_in_metadata_thread( - _filter_tools_by_current_user_permission, tools - ) + return await _filter_visible_tools_fail_open(tools) def _make_call_tool(self) -> Any: """Build the normalized ``call_tool`` proxy for regex search.""" @@ -695,9 +715,7 @@ def _create_search_transform( # noqa: C901 tools = await super()._get_visible_tools(ctx) # Permission lookups need a metadata connection; see # RBACToolVisibilityMiddleware.on_list_tools. - return await run_in_metadata_thread( - _filter_tools_by_current_user_permission, tools - ) + return await _filter_visible_tools_fail_open(tools) def _make_call_tool(self) -> Any: """Build the normalized ``call_tool`` proxy for BM25 search.""" diff --git a/superset/mcp_service/worker.py b/superset/mcp_service/worker.py index ea7df0f6e13..06835a76f33 100644 --- a/superset/mcp_service/worker.py +++ b/superset/mcp_service/worker.py @@ -166,15 +166,26 @@ class WorkerPool: try: future = self.cancellations.submit(fn) except RuntimeError: + # The cancellations executor is shut down; report failure to + # dispatch rather than raise. self.cancel_slots.release() return False + except BaseException: + self.cancel_slots.release() + raise def completed(_: Future[None]) -> None: """Release cancellation capacity before admitting another tool call.""" self.cancel_slots.release() finished() - future.add_done_callback(completed) + try: + future.add_done_callback(completed) + except BaseException: + # The permit is otherwise never released: completed() runs only if + # the callback was registered, and a raise here means it was not. + self.cancel_slots.release() + raise return True @@ -235,18 +246,44 @@ def tool_worker_count(app: Flask) -> int: def metadata_tool_worker_count(app: Flask) -> int: - """Admit metadata-only calls independently of warehouse capacity. + """Admit metadata-only calls up to whatever the warehouse bound leaves spare. They hold a metadata connection only for their own short metadata queries, - never across warehouse I/O, so they cannot keep the connections reserved - for warehouse calls and their cancellation from cycling. They only wait - briefly for a connection, on a worker thread rather than the event loop. + never across warehouse I/O, so they do not need the two-connections-per-call + margin ``tool_worker_count`` reserves. But they draw from the same pool: that + reservation (``2 * workers + 1``) is sized to use the *entire* bounded pool + by default, so admitting metadata-only calls beyond what is left over would + let them, in aggregate, check out more metadata connections than the pool + has -- the exact exhaustion this module exists to prevent, just moved to a + different set of tools. With an unbounded pool (``_metadata_pool_capacity`` + returns ``None``) there is nothing to divide, so the configured/default + count is used as-is. """ configured = app.config.get("MCP_METADATA_TOOL_WORKERS") - count = DEFAULT_METADATA_TOOL_WORKERS if configured is None else configured - if count < 0: + requested = DEFAULT_METADATA_TOOL_WORKERS if configured is None else configured + if requested < 0: raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative") - return count + capacity = _metadata_pool_capacity(app) + if capacity is None: + return requested + reserved_for_warehouse = 2 * tool_worker_count(app) + 1 + if requested > (remaining := max(capacity - reserved_for_warehouse, 0)): + if configured is not None: + logger.warning( + "MCP_METADATA_TOOL_WORKERS=%s would need %s metadata database " + "connections on top of the %s reserved for warehouse-capable " + "tools, but the pool only allows %s; admitting %s concurrent " + "metadata-only calls instead. Raise the pool's " + "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS, or lower " + "MCP_TOOL_WORKERS, to admit more.", + configured, + requested, + reserved_for_warehouse, + capacity, + remaining, + ) + return remaining + return requested def _get_pool(app: Flask) -> WorkerPool: @@ -367,12 +404,20 @@ class TransportContext: async def forward(*args: Any, **kwargs: Any) -> Any: self.call.check() future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop) + wrapped = asyncio.wrap_future(future) try: result = await asyncio.wait_for( - asyncio.wrap_future(future), + wrapped, max(0, self.call.deadline - time.monotonic()), ) except TimeoutError: + if wrapped.done() and not wrapped.cancelled(): + # wait_for cancels the awaited future on its own timeout + # (no shield here), so a TimeoutError that leaves it + # already done rather than cancelled was genuinely + # raised by the forwarded call itself. Propagate that + # one as itself, not relabeled as our deadline. + raise raise WorkerDeadlineExceeded() from None self.call.check() return result @@ -466,6 +511,41 @@ async def run_in_metadata_thread( return await asyncio.to_thread(execute) +def _handle_worker_timeout_or_cancel( + exc: BaseException, + wrapped: asyncio.Future[Any], + call: WorkerCall, + seconds: float, +) -> None: + """Decide how a run_in_worker deadline or cancellation should surface. + + Always raises; the caller's except block delegates here instead of + inlining the decision, to keep run_in_worker's own branching simple. + """ + call.abandon() + # Retrieve late exceptions, including worker CancelledError, without + # retaining a task on the transport loop after the request has ended. + wrapped.add_done_callback( + lambda done: None if done.cancelled() else done.exception() + ) + if isinstance(exc, asyncio.CancelledError): + raise exc + if isinstance(exc, TimeoutError) and wrapped.done(): + # wait_for raises this same exception type both on its own deadline + # and when the shielded future already finished with a TimeoutError + # of its own (e.g. a driver-level socket timeout). Shielding means + # our deadline firing leaves the future still pending; only a + # TimeoutError genuinely raised by the tool/driver leaves it already + # done. Propagate that one as itself rather than relabeling it as + # our deadline. + raise exc + raise ToolError( + f"MCP tool timed out after {seconds:g} seconds. " + "Warehouse cancellation was requested where supported. " + f"Call id: {call.call_id}" + ) from None + + async def run_in_worker( fn: Callable[..., Coroutine[Any, Any, Any]], args: tuple[Any, ...], @@ -548,27 +628,18 @@ async def run_in_worker( call.check() return result - future = pool.submit(lambda: context.run(execute), call.finished, call.admitted) - wrapped = asyncio.wrap_future(future) try: - # Shield the future: cancellation must not release its pool slot early. - return await asyncio.wait_for( - asyncio.shield(wrapped), timeout=max(0, call.deadline - time.monotonic()) - ) - except (TimeoutError, WorkerDeadlineExceeded, asyncio.CancelledError) as exc: - call.abandon() - # Retrieve late exceptions, including worker CancelledError, without - # retaining a task on the transport loop after the request has ended. - wrapped.add_done_callback( - lambda done: None if done.cancelled() else done.exception() - ) - if isinstance(exc, asyncio.CancelledError): - raise - raise ToolError( - f"MCP tool timed out after {seconds:g} seconds. " - "Warehouse cancellation was requested where supported. " - f"Call id: {call.call_id}" - ) from None + future = pool.submit(lambda: context.run(execute), call.finished, call.admitted) + wrapped = asyncio.wrap_future(future) + try: + # Shield the future: cancellation must not release its pool slot + # early. + return await asyncio.wait_for( + asyncio.shield(wrapped), + timeout=max(0, call.deadline - time.monotonic()), + ) + except (TimeoutError, WorkerDeadlineExceeded, asyncio.CancelledError) as exc: + _handle_worker_timeout_or_cancel(exc, wrapped, call, seconds) finally: from superset.mcp_service.auth import _mcp_user_id_var @@ -597,9 +668,16 @@ class QueryCancellation: self.context = copy_context() def capture(self) -> None: - """Capture an engine handle without making unsupported drivers fail.""" + """Capture an engine handle without making unsupported drivers fail. + + Cleared up front: a multi-statement call re-captures per statement + (see ``refresh``), and a failed re-capture must not leave register() + republishing a stale handle for a statement that has already + finished. + """ from superset.tasks.query_cancel import capture_cancel_query_id + self.cancel_id = None try: self.cancel_id = capture_cancel_query_id(self.database, self.cursor) except Exception: diff --git a/tests/unit_tests/mcp_service/test_middleware.py b/tests/unit_tests/mcp_service/test_middleware.py index 47c3f054757..d6decb56e53 100644 --- a/tests/unit_tests/mcp_service/test_middleware.py +++ b/tests/unit_tests/mcp_service/test_middleware.py @@ -2426,6 +2426,27 @@ class TestRBACToolVisibilityMiddleware: assert result == tools + @pytest.mark.asyncio + async def test_fails_open_when_metadata_thread_setup_raises(self) -> None: + """Returns all tools when run_in_metadata_thread's own setup fails. + + Unlike test_fails_open_on_exception, this failure happens before + _visible_tools ever runs -- e.g. the metadata pool is exhausted while + run_in_metadata_thread reloads the caller's ORM user -- so + _visible_tools's own try/except never gets a chance to fail open. + """ + tools = [self._make_tool("list_charts"), self._make_tool("generate_chart")] + call_next = AsyncMock(return_value=tools) + middleware = RBACToolVisibilityMiddleware() + + with patch( + "superset.mcp_service.middleware.run_in_metadata_thread", + side_effect=RuntimeError("metadata pool exhausted"), + ): + result = await middleware.on_list_tools(MagicMock(), call_next) + + assert result == tools + @pytest.mark.asyncio async def test_fails_open_when_user_is_none(self, app) -> None: """Returns all tools when get_user_from_request returns None.""" diff --git a/tests/unit_tests/mcp_service/test_worker.py b/tests/unit_tests/mcp_service/test_worker.py index 35a008ecafe..a308f838fdf 100644 --- a/tests/unit_tests/mcp_service/test_worker.py +++ b/tests/unit_tests/mcp_service/test_worker.py @@ -24,7 +24,7 @@ from collections.abc import Iterator from contextlib import contextmanager, ExitStack from contextvars import ContextVar from typing import Any -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest from fastmcp import Client, FastMCP @@ -380,6 +380,58 @@ async def test_failed_warehouse_discards_metadata_session(app: Any) -> None: pool.cancellations.shutdown() +def test_query_cancellation_capture_resets_stale_id_on_failed_recapture() -> None: + """A failed re-capture must not leave register() republishing a stale id + from an earlier, already-finished statement in a multi-statement call. + """ + from superset.mcp_service.worker import QueryCancellation + + call = MagicMock() + call.expired.is_set.return_value = False + database = MagicMock() + database.db_engine_spec.has_implicit_cancel.return_value = False + cancellation = QueryCancellation(call, database, Mock(), None, None) + + with patch( + "superset.tasks.query_cancel.capture_cancel_query_id", + return_value="statement-1-id", + ): + cancellation.capture() + assert cancellation.cancel_id == "statement-1-id" + cancellation.register() + assert call.cancel_query is not None + + with patch( + "superset.tasks.query_cancel.capture_cancel_query_id", + side_effect=RuntimeError("driver refused a second handle"), + ): + cancellation.capture() + assert cancellation.cancel_id is None + cancellation.register() + # No usable handle for this statement: cancellation must go unavailable, + # not silently keep dispatching against the finished first statement. + assert call.cancel_query is None + + +def test_worker_pool_cancel_releases_slot_on_unexpected_submit_error() -> None: + """Any failure to dispatch cancellation I/O must release its permit, not + just the already-handled RuntimeError-on-shutdown case.""" + from superset.mcp_service.worker import WorkerPool + + pool = WorkerPool(1) + try: + with patch.object(pool.cancellations, "submit", side_effect=ValueError("boom")): + with pytest.raises(ValueError, match="boom"): + pool.cancel(Mock(), Mock()) + # The permit was released despite the unexpected exception type, so a + # later cancellation can still be admitted. + assert pool.cancel_slots.acquire(blocking=False) + pool.cancel_slots.release() + finally: + pool.executor.shutdown() + pool.cancellations.shutdown() + + @pytest.mark.asyncio async def test_sqlalchemy_warehouse_queries_register_cancellation(app: Any) -> None: """Metadata discovery also uses the engine hooks, not just chart/SQL Lab SQL.""" diff --git a/tests/unit_tests/mcp_service/test_worker_metadata_pool.py b/tests/unit_tests/mcp_service/test_worker_metadata_pool.py index bf7ba6ebf47..57f5646e75a 100644 --- a/tests/unit_tests/mcp_service/test_worker_metadata_pool.py +++ b/tests/unit_tests/mcp_service/test_worker_metadata_pool.py @@ -640,7 +640,8 @@ async def test_execute_sql_reports_missing_table_as_query_error( def test_metadata_only_tools_have_their_own_bound( app: Any, configured: int | None, expected: int ) -> None: - """The metadata-only bound does not depend on warehouse reservations.""" + """With an unbounded pool there is nothing to divide, so the configured or + default count is used as-is.""" from superset.mcp_service.worker import metadata_tool_worker_count with patch.dict(app.config, {"MCP_METADATA_TOOL_WORKERS": configured}): @@ -652,6 +653,54 @@ def test_metadata_only_tools_have_their_own_bound( metadata_tool_worker_count(app) [email protected]( + ("pool_options", "tool_workers", "configured", "expected"), + [ + # SQLAlchemy's default QueuePool lends 5 + 10 = 15 connections. + # tool_worker_count(None) picks 7 (2*7+1 == 15): nothing left over. + ({}, None, None, 0), + ({}, None, 3, 0), + # A smaller warehouse bound leaves room for metadata-only calls. + ({}, 2, None, 10), + ({}, 2, 6, 6), + ({}, 2, 20, 10), + # A generous pool still caps the default at what's left once + # tool_worker_count's own default (capped at 16) is reserved for. + ({"pool_size": 20, "max_overflow": 20}, None, None, 7), + ], +) +def test_metadata_only_bound_is_capped_by_remaining_pool_capacity( + app: Any, + tmp_path: Path, + pool_options: dict[str, Any], + tool_workers: int | None, + configured: int | None, + expected: int, +) -> None: + """Metadata-only calls cannot collectively outnumber what the pool, net of + the warehouse reservation, can actually serve -- the same connection + exhaustion this module exists to prevent, just moved to a different set + of tools.""" + from superset.mcp_service.worker import metadata_tool_worker_count + + engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}", **pool_options) + previous = { + key: app.config.get(key) + for key in ("MCP_TOOL_WORKERS", "MCP_METADATA_TOOL_WORKERS") + } + app.config["MCP_TOOL_WORKERS"] = tool_workers + app.config["MCP_METADATA_TOOL_WORKERS"] = configured + try: + with patch.object( + type(db), "engine", new_callable=PropertyMock, return_value=engine + ): + assert metadata_tool_worker_count(app) == expected + finally: + app.config["MCP_TOOL_WORKERS"] = previous["MCP_TOOL_WORKERS"] + app.config["MCP_METADATA_TOOL_WORKERS"] = previous["MCP_METADATA_TOOL_WORKERS"] + engine.dispose() + + def test_metadata_only_tools_are_registered_tools() -> None: """A renamed tool must not silently fall out of (or into) the fast bound.""" from superset.mcp_service.app import mcp diff --git a/tests/unit_tests/sql/execution/test_cancellation.py b/tests/unit_tests/sql/execution/test_cancellation.py index 2f65268813f..51733fc856b 100644 --- a/tests/unit_tests/sql/execution/test_cancellation.py +++ b/tests/unit_tests/sql/execution/test_cancellation.py @@ -100,6 +100,32 @@ def test_engine_cursor_scope_lifecycle(fail: bool) -> None: metadata.dispose() +def test_cancellable_engine_admits_before_first_connection() -> None: + """A connect-event prequery is admitted before cancellable_engine yields. + + ``before_warehouse_access`` is the only hook a raw DBAPI connect-event + prequery can reach -- it runs before SQLAlchemy's cursor-execute events + exist to register with -- so cancellable_engine must call it unconditionally + on entry, not only once a cursor is later registered. + """ + admit = Mock() + engine = create_engine("sqlite://") + try: + with ExitStack() as stack: + stack.callback( + cancellation.cursor_scope.reset, cancellation.cursor_scope.set(Mock()) + ) + stack.callback( + cancellation.before_warehouse_access.reset, + cancellation.before_warehouse_access.set(admit), + ) + admit.assert_not_called() + with cancellation.cancellable_engine(Mock(), engine, None, None): + admit.assert_called_once_with() + finally: + engine.dispose() + + def test_cursor_and_engine_without_owner() -> None: """Web/Celery execution does not install a cursor scope or call hooks.""" database = Mock()
