aminghadersohi commented on code in PR #44581: URL: https://github.com/apache/superset/pull/44581#discussion_r4150911741
########## superset/sql/execution/cancellation.py: ########## @@ -0,0 +1,170 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Optional execution-owner hooks around a live warehouse cursor. + +Ordinary web and Celery execution are unchanged. An in-process execution owner +can register cancellation while keeping its async runtime out of model code. +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, contextmanager +from contextvars import ContextVar +from typing import Any, Callable, Iterator, TYPE_CHECKING + +from sqlalchemy import event +from sqlalchemy.engine import Engine + +if TYPE_CHECKING: + from superset.models.core import Database + +cursor_scope: ContextVar[ + Callable[[Database, Any, str | None, str | None], AbstractContextManager[None]] + | None +] = ContextVar("warehouse_cursor_scope", default=None) + + +@contextmanager +def cancellable_cursor( + database: Database, + cursor: Any, + catalog: str | None = None, + schema: str | None = None, +) -> Iterator[None]: + """Let the execution owner capture cancellation before blocking DBAPI I/O.""" + scope = cursor_scope.get() + if scope is None: + yield + else: + with scope(database, cursor, catalog, schema): + yield + + +before_warehouse_access: ContextVar[Callable[[], None] | None] = ContextVar( + "warehouse_before_access", default=None +) +check_deadline: ContextVar[Callable[[], None] | None] = ContextVar( + "warehouse_check_deadline", default=None +) +after_execute: ContextVar[Callable[[], None] | None] = ContextVar( + "warehouse_after_execute", default=None +) + + +def check_query_deadline() -> None: + """Stop before metadata access or another statement after warehouse I/O.""" + if check := check_deadline.get(): + check() + + +def query_executed() -> None: + """Refresh cancellation handles exposed only after driver execution.""" + if refresh := after_execute.get(): + refresh() + + +_engine_scope: ContextVar[tuple[Database, Engine, str | None, str | None] | None] = ( + ContextVar("warehouse_engine_scope", default=None) +) + + +@contextmanager +def without_execution_hooks() -> Iterator[None]: + """Keep cancellation I/O independent of the query it is cancelling. + + Preserve unrelated context, such as authentication and tenant routing, while + preventing prequeries from checking the abandoned deadline or registering + cancellation recursively. + """ + cursor_token = cursor_scope.set(None) + access_token = before_warehouse_access.set(None) + deadline_token = check_deadline.set(None) + execute_token = after_execute.set(None) + engine_token = _engine_scope.set(None) + try: + yield + finally: + _engine_scope.reset(engine_token) + after_execute.reset(execute_token) + check_deadline.reset(deadline_token) + before_warehouse_access.reset(access_token) + cursor_scope.reset(cursor_token) + + +@contextmanager +def cancellable_engine( + database: Database, engine: Engine, catalog: str | None, schema: str | None +) -> Iterator[None]: + """Cover SQLAlchemy warehouse statements, including metadata discovery. + + Listeners are installed once, never added/removed on shared cached engines. + Matching the engine excludes Superset metadata queries in the same context. + """ + if cursor_scope.get() is None: + yield + return + if admit := before_warehouse_access.get(): + # Before any connection is opened, including prequeries on connect. + admit() Review Comment: Added your test as `test_engine_admits_before_opening_a_connection` in `tests/unit_tests/sql/execution/test_cancellation.py` (91a3a66b0e). Locally, `pytest --cov=superset/sql/ ./tests/unit_tests/sql/ --cov-fail-under=100` now gives `cancellation.py` 100% and `TOTAL ... 100%` ("Required test coverage of 100% reached. Total coverage: 100.00%"), with 2,012 passed. ########## superset/mcp_service/worker.py: ########## @@ -0,0 +1,700 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Bounded, deadline-aware execution of MCP tools outside the transport loop. + +The tool's coroutine runs on a worker-owned loop because its database APIs are +synchronous. Only transport notifications are marshalled back to the server +loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting +request: a timed-out DBAPI call cannot outlive and reuse a torn-down session. +""" + +from __future__ import annotations + +import asyncio +import functools +import logging +import threading +import time +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from contextvars import ContextVar, copy_context +from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar +from weakref import WeakKeyDictionary + +from fastmcp.exceptions import ToolError +from flask import current_app, g, has_app_context, has_request_context +from sqlalchemy import inspect as sa_inspect +from sqlalchemy.orm.state import InstanceState +from sqlalchemy.pool import QueuePool + +from superset.mcp_service.session_scope import _mcp_session_token + +if TYPE_CHECKING: + from flask import Flask + + from superset.models.core import Database + +logger = logging.getLogger(__name__) +_active_call: ContextVar[WorkerCall | None] = ContextVar( + "mcp_worker_call", default=None +) +_metadata_context_owned: ContextVar[bool] = ContextVar( + "mcp_metadata_context_owned", default=False +) +_P = ParamSpec("_P") +_T = TypeVar("_T") +_pools_lock = threading.Lock() +_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary() + + +class WorkerDeadlineExceeded(BaseException): + """Stop abandoned work without being swallowed by tool error handlers.""" + + +BUSY_MESSAGE = "MCP server busy: all tool workers are occupied. Retry later." + +# Tools that only read Superset's metadata database. They are admitted under a +# separate, larger bound so they keep answering while slow warehouse queries +# hold every warehouse slot. Anything that can reach a warehouse, a semantic +# layer or a screenshot service stays in the warehouse bound. A listed tool +# that does reach a warehouse must still take a warehouse slot first (see +# WorkerCall.admit_warehouse), so a wrong entry cannot overdraw the pool. +METADATA_ONLY_TOOLS = frozenset( + { + "find_users", + "get_annotation_layer_info", + "get_chart_info", + "get_chart_type_schema", + "get_dashboard_info", + "get_database_info", + "get_dataset_info", + "get_instance_info", + "get_layer_annotation_info", + "get_query_info", + "get_report_info", + "get_rls_filter_info", + "get_role_info", + "get_saved_query_info", + "get_schema", + "get_tag_info", + "get_task_info", + "get_theme_info", + "get_user_info", + "health_check", + "list_annotation_layers", + "list_charts", + "list_dashboards", + "list_databases", + "list_datasets", + "list_layer_annotations", + "list_queries", + "list_reports", + "list_rls_filters", + "list_roles", + "list_saved_queries", + "list_tags", + "list_tasks", + "list_themes", + "list_users", + } +) + + +class WorkerPool: + """Bound submissions, including abandoned work; never queue behind a query.""" + + def __init__(self, size: int, metadata_size: int = 0) -> None: + if size < 1: + raise ValueError("MCP_TOOL_WORKERS must be positive") + if metadata_size < 0: + raise ValueError("MCP_METADATA_TOOL_WORKERS must not be negative") + self.slots = threading.BoundedSemaphore(size) + self.metadata_slots = ( + threading.BoundedSemaphore(metadata_size) if metadata_size else None + ) + self.executor = ThreadPoolExecutor( + size + metadata_size, thread_name_prefix="mcp-tool" + ) + # Cancellation must not wait behind the warehouse work it is cancelling. + # Only calls holding a warehouse slot register cancellation. + self.cancel_slots = threading.BoundedSemaphore(size) + self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel") + + def admission(self, metadata_only: bool) -> threading.BoundedSemaphore: + """Choose the bound a call is admitted under.""" + if metadata_only and self.metadata_slots is not None: + return self.metadata_slots + return self.slots + + def submit( + self, + fn: Callable[[], Any], + finished: Callable[[], None], + slots: threading.BoundedSemaphore | None = None, + ) -> Future[Any]: + """Admit immediately or report overload without retaining a queued call.""" + slots = self.slots if slots is None else slots + if not slots.acquire(blocking=False): + raise ToolError(BUSY_MESSAGE) + try: + future = self.executor.submit(fn) + except BaseException: + slots.release() + raise + future.add_done_callback(lambda _: finished()) + return future + + def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool: + """Bound cancellation I/O separately so it cannot delay the caller.""" + if not self.cancel_slots.acquire(blocking=False): + return False + try: + future = self.cancellations.submit(fn) + except RuntimeError: + self.cancel_slots.release() + return False + + def completed(_: Future[None]) -> None: + """Release cancellation capacity before admitting another tool call.""" + self.cancel_slots.release() + finished() + + future.add_done_callback(completed) + return True + + +DEFAULT_TOOL_WORKERS = 16 +DEFAULT_METADATA_TOOL_WORKERS = 16 + + +def _metadata_pool_capacity(app: Flask) -> int | None: + """Return how many metadata connections can be checked out at once. + + ``None`` means a checkout never waits for another holder to return one + (e.g. ``NullPool``, per-thread pools, or unlimited overflow). + """ + from superset import db + + with app.app_context(): + pool = db.engine.pool + if not isinstance(pool, QueuePool): + return None + # SQLAlchemy has no public accessor for the configured overflow limit. + max_overflow = pool._max_overflow # pylint: disable=protected-access + return None if max_overflow < 0 else pool.size() + max_overflow + + +def tool_worker_count(app: Flask) -> int: + """Admit only as many calls as the metadata pool can always serve. + + An admitted call can hold one metadata connection for the whole of its + warehouse I/O, and its cancellation needs another. ``2 * workers + 1`` + connections therefore always leave one that is only held by short metadata + lookups, so cancellation and transport-side lookups (tools/list filtering, + audit logging) never wait for a warehouse query to end on its own. + """ + configured = app.config.get("MCP_TOOL_WORKERS") + capacity = _metadata_pool_capacity(app) + if capacity is None: + return DEFAULT_TOOL_WORKERS if configured is None else configured + limit = (capacity - 1) // 2 + if limit < 1: + raise ValueError( + f"The metadata database pool allows {capacity} connections; " + "MCP tool execution needs at least 3" + ) + if configured is None: + return min(DEFAULT_TOOL_WORKERS, limit) + if configured > limit: + logger.warning( + "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the " + "pool allows %s; admitting %s concurrent tool calls. Raise the pool's " + "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.", + configured, + 2 * configured + 1, + capacity, + limit, + ) + return limit + return configured + + +def metadata_tool_worker_count(app: Flask) -> int: + """Admit metadata-only calls independently of warehouse capacity. + + 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. + """ + configured = app.config.get("MCP_METADATA_TOOL_WORKERS") + count = DEFAULT_METADATA_TOOL_WORKERS if configured is None else configured Review Comment: Clamped in c9bc7895ae. Both bounds now come out of one budget, checked when the pool is created at startup (`admission_counts`): `2 * MCP_TOOL_WORKERS + MCP_METADATA_TOOL_WORKERS + 1 <= pool_size + max_overflow` - The `+ 1` is now an actual bound. Transport-side metadata I/O used the unbounded default executor; it now runs on one dedicated thread. That covers tools/list filtering, audit writes, error hooks, and `CompositeTokenVerifier`'s API-key lookup. - With the default 5 + 10 pool the defaults are 4 warehouse-capable and 6 metadata-only calls, down from 7 and 16, so the default warehouse concurrency drops from 7 to 4. Explicit values above the budget are reduced with a startup warning. Docs and `mcp_config.py` describe the budget. - `test_saturated_default_pool_never_waits_for_a_connection` runs on a real 5 + 10 `QueuePool` with `pool_timeout=1`. It holds all 4 warehouse slots, 4 cancellation connections, all 6 metadata-only slots and the transport thread at once (15 checked out). Every holder gets its connection and both bounds refuse the next call. Re-run with one metadata slot over the budget, it fails with `BrokenBarrierError`. - `test_tool_workers_leave_metadata_connections_for_cancellation` checks the split and the inequality across 11 pool and config combinations. - d8f971db50 fixes a side effect I hit while doing this. Creating the pool lazily from transport-side code popped an app context, and that removed the caller's scoped session. Covered by `test_lazy_pool_creation_keeps_callers_session`, which fails without the fix. ########## superset/mcp_service/worker.py: ########## @@ -0,0 +1,588 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Bounded, deadline-aware execution of MCP tools outside the transport loop. + +The tool's coroutine runs on a worker-owned loop because its database APIs are +synchronous. Only transport notifications are marshalled back to the server +loop. Flask and SQLAlchemy lifetimes belong to the worker, not the waiting +request: a timed-out DBAPI call cannot outlive and reuse a torn-down session. +""" + +from __future__ import annotations + +import asyncio +import functools +import logging +import threading +import time +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from contextvars import ContextVar, copy_context +from typing import Any, Callable, Coroutine, Iterator, ParamSpec, TYPE_CHECKING, TypeVar +from weakref import WeakKeyDictionary + +from fastmcp.exceptions import ToolError +from flask import current_app, g, has_app_context, has_request_context +from sqlalchemy import inspect as sa_inspect +from sqlalchemy.orm.state import InstanceState +from sqlalchemy.pool import QueuePool + +from superset.mcp_service.session_scope import _mcp_session_token + +if TYPE_CHECKING: + from flask import Flask + + from superset.models.core import Database + +logger = logging.getLogger(__name__) +_active_call: ContextVar[WorkerCall | None] = ContextVar( + "mcp_worker_call", default=None +) +_metadata_context_owned: ContextVar[bool] = ContextVar( + "mcp_metadata_context_owned", default=False +) +_P = ParamSpec("_P") +_T = TypeVar("_T") +_pools_lock = threading.Lock() +_pools: WeakKeyDictionary[Flask, WorkerPool] = WeakKeyDictionary() + + +class WorkerDeadlineExceeded(BaseException): + """Stop abandoned work without being swallowed by tool error handlers.""" + + +class WorkerPool: + """Bound submissions, including abandoned work; never queue behind a query.""" + + def __init__(self, size: int) -> None: + if size < 1: + raise ValueError("MCP_TOOL_WORKERS must be positive") + self.slots = threading.BoundedSemaphore(size) + self.executor = ThreadPoolExecutor(size, thread_name_prefix="mcp-tool") + # Cancellation must not wait behind the warehouse work it is cancelling. + self.cancel_slots = threading.BoundedSemaphore(size) + self.cancellations = ThreadPoolExecutor(size, thread_name_prefix="mcp-cancel") + + def submit( + self, fn: Callable[[], Any], finished: Callable[[], None] + ) -> Future[Any]: + """Admit immediately or report overload without retaining a queued call.""" + if not self.slots.acquire(blocking=False): + raise ToolError( + "MCP server busy: all tool workers are occupied. Retry later." + ) + try: + future = self.executor.submit(fn) + except BaseException: + self.slots.release() + raise + future.add_done_callback(lambda _: finished()) + return future + + def cancel(self, fn: Callable[[], None], finished: Callable[[], None]) -> bool: + """Bound cancellation I/O separately so it cannot delay the caller.""" + if not self.cancel_slots.acquire(blocking=False): + return False + try: + future = self.cancellations.submit(fn) + except RuntimeError: + self.cancel_slots.release() + return False + + def completed(_: Future[None]) -> None: + """Release cancellation capacity before admitting another tool call.""" + self.cancel_slots.release() + finished() + + future.add_done_callback(completed) + return True + + +DEFAULT_TOOL_WORKERS = 16 + + +def _metadata_pool_capacity(app: Flask) -> int | None: + """Return how many metadata connections can be checked out at once. + + ``None`` means a checkout never waits for another holder to return one + (e.g. ``NullPool``, per-thread pools, or unlimited overflow). + """ + from superset import db + + with app.app_context(): + pool = db.engine.pool + if not isinstance(pool, QueuePool): + return None + # SQLAlchemy has no public accessor for the configured overflow limit. + max_overflow = pool._max_overflow # pylint: disable=protected-access + return None if max_overflow < 0 else pool.size() + max_overflow + + +def tool_worker_count(app: Flask) -> int: + """Admit only as many calls as the metadata pool can always serve. + + An admitted call can hold one metadata connection for the whole of its + warehouse I/O, and its cancellation needs another. ``2 * workers + 1`` + connections therefore always leave one that is only held by short metadata + lookups, so cancellation and transport-side lookups (tools/list filtering, + audit logging) never wait for a warehouse query to end on its own. + """ + configured = app.config.get("MCP_TOOL_WORKERS") + capacity = _metadata_pool_capacity(app) + if capacity is None: + return DEFAULT_TOOL_WORKERS if configured is None else configured + limit = (capacity - 1) // 2 + if limit < 1: + raise ValueError( + f"The metadata database pool allows {capacity} connections; " + "MCP tool execution needs at least 3" + ) + if configured is None: + return min(DEFAULT_TOOL_WORKERS, limit) + if configured > limit: + logger.warning( + "MCP_TOOL_WORKERS=%s needs %s metadata database connections, but the " + "pool allows %s; admitting %s concurrent tool calls. Raise the pool's " + "pool_size/max_overflow in SQLALCHEMY_ENGINE_OPTIONS to admit more.", + configured, + 2 * configured + 1, + capacity, + limit, + ) + return limit + return configured + + +def _get_pool(app: Flask) -> WorkerPool: + """Lazily create a pool for this application, without import-time threads.""" + with _pools_lock: + if app not in _pools: + size = tool_worker_count(app) + logger.info("MCP tool calls admitted concurrently: %s", size) + _pools[app] = WorkerPool(size) + return _pools[app] + + +class WorkerCall: + """Thread-safe deadline and cancellation registration for one tool call.""" + + def __init__(self, app: Flask, pool: WorkerPool, seconds: float) -> None: + self.app = app + self.pool = pool + self.loop = asyncio.get_running_loop() + self.seconds = seconds + self.deadline = time.monotonic() + seconds + from superset.mcp_service.middleware import _mcp_call_id_var + + self.call_id = _mcp_call_id_var.get() or uuid.uuid4().hex + self.expired = threading.Event() + self.lock = threading.RLock() + self.pending = 1 + self.cancel_query: Callable[[], None] | None = None + self.cancel_dispatched = False + self.user_id: int | None = None + + def check(self) -> None: + """Prevent an abandoned tool from starting more work or mutations.""" + if self.expired.is_set() or time.monotonic() >= self.deadline: + # BaseException deliberately bypasses tools' broad Exception handlers. + raise WorkerDeadlineExceeded() + + def dispatch_cancel(self) -> bool: + """Dispatch at most once for the active cursor, including late handles.""" + with self.lock: + if self.cancel_query is None or self.cancel_dispatched: + return self.cancel_dispatched + self.pending += 1 + self.cancel_dispatched = self.pool.cancel(self.cancel_query, self.finished) + if not self.cancel_dispatched: + self.pending -= 1 + return self.cancel_dispatched + + def finished(self) -> None: + """Retain admission until both query and cancellation I/O have ended. + + A stuck cancellation therefore cannot consume the cancellation capacity + needed by newly admitted queries: it retains its original tool slot. + """ + with self.lock: + self.pending -= 1 + if self.pending == 0: + self.pool.slots.release() + + def abandon(self) -> None: + """Signal abandonment and dispatch cancellation without blocking asyncio.""" + self.expired.set() + dispatched = self.dispatch_cancel() + logger.warning( + "MCP call %s exceeded its deadline or disconnected; " + "cancellation dispatched=%s; worker slot retained until completion", + self.call_id, + dispatched, + ) + + +class TransportContext: + """Keep FastMCP's transport-bound async methods on their owning event loop.""" + + def __init__( + self, context: Any, loop: asyncio.AbstractEventLoop, call: WorkerCall + ) -> None: + self.context = context + self.loop = loop + self.call = call + + def __getattr__(self, name: str) -> Any: + value = getattr(self.context, name) + if not asyncio.iscoroutinefunction(value): + return value + + @functools.wraps(value) + async def forward(*args: Any, **kwargs: Any) -> Any: + self.call.check() + future = asyncio.run_coroutine_threadsafe(value(*args, **kwargs), self.loop) + try: + result = await asyncio.wait_for( + asyncio.wrap_future(future), + max(0, self.call.deadline - time.monotonic()), + ) + except TimeoutError: + raise WorkerDeadlineExceeded() from None + self.call.check() + return result + + return forward + + +@contextmanager +def _worker_context(app: Flask) -> Iterator[None]: + """Give a worker an independent app context and scoped metadata session.""" + from superset.mcp_service.auth import _remove_session_safe + + token = _mcp_session_token.set(object()) + try: + with app.app_context(): + try: + # Clean this scope before loading any ORM user into it. + _remove_session_safe() + yield + finally: + _remove_session_safe() + finally: + _mcp_session_token.reset(token) + + +def get_context_user_id() -> int | None: + """Read the caller's identity without refreshing an expired ORM instance.""" + if not has_app_context(): + return None + user = getattr(g, "user", None) + state = sa_inspect(user, raiseerr=False) + if isinstance(state, InstanceState): + return state.identity[0] if state.identity else None + return getattr(user, "id", None) + + +async def run_in_metadata_thread( + fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs +) -> _T: + """Run transport metadata I/O with thread-owned Flask and session lifetimes. + + ``to_thread`` copies contextvars, including Flask contexts and MCP session + tokens. Replace those owners, reload ORM users, and retain only request and + routing data. Cleanup belongs to the thread even if its awaiter disconnects. + This uses the metadata executor rather than admitted tool workers, which + may themselves be waiting for transport notifications. + """ + from contextlib import nullcontext + + from flask.globals import _cv_request + + if has_app_context(): + app = current_app._get_current_object() + snapshot = dict(vars(g._get_current_object())) + else: + from superset.mcp_service.flask_singleton import get_flask_app + + app = get_flask_app() + snapshot = {} + user = snapshot.pop("user", None) + state = sa_inspect(user, raiseerr=False) + is_orm_user = isinstance(state, InstanceState) + user_id = get_context_user_id() if is_orm_user else None + request_context = _cv_request.get(None) + request_copy = request_context.copy() if request_context is not None else None + + def execute() -> _T: + """Own teardown rather than handing a live session back to asyncio.""" + from superset import db, security_manager + from superset.sql.execution.cancellation import without_execution_hooks + + active_token = _active_call.set(None) + owner_token = _metadata_context_owned.set(True) + try: + with without_execution_hooks(), _worker_context(app): + vars(g._get_current_object()).update(snapshot) + with request_copy if request_copy is not None else nullcontext(): + if is_orm_user: + g.user = ( + db.session.get(security_manager.user_model, user_id) + if user_id is not None + else None + ) + elif user is not None: + g.user = user + return fn(*args, **kwargs) + finally: + _metadata_context_owned.reset(owner_token) + _active_call.reset(active_token) + + return await asyncio.to_thread(execute) + + +async def run_in_worker( + fn: Callable[..., Coroutine[Any, Any, Any]], + args: tuple[Any, ...], + kwargs: dict[str, Any], + seconds: float, +) -> Any: + """Run a complete tool lifecycle in a bounded, independently owned worker.""" + if active := _active_call.get(): + # Composed tools share the outer deadline and worker/session ownership. + # A nested auth hook can inject the original transport Context again. + nested_kwargs = dict(kwargs) + if "ctx" in nested_kwargs and not isinstance( + nested_kwargs["ctx"], TransportContext + ): + nested_kwargs["ctx"] = TransportContext( + nested_kwargs["ctx"], active.loop, active + ) + active.check() + return await fn(*args, **nested_kwargs) + + if has_app_context(): + app = current_app._get_current_object() + else: + from superset.mcp_service.flask_singleton import get_flask_app + + app = get_flask_app() + pool = _get_pool(app) + call = WorkerCall(app, pool, seconds) + loop = asyncio.get_running_loop() + worker_kwargs = dict(kwargs) + if "ctx" in worker_kwargs: + worker_kwargs["ctx"] = TransportContext(worker_kwargs["ctx"], loop, call) + # Copy contextvars (token, tenant routing), but never share Flask g or a + # metadata Session. Request-backed middleware's user is reloaded by id. + from flask.globals import _cv_request + + request_context = _cv_request.get(None) + request_copy = request_context.copy() if request_context is not None else None + globals_snapshot = dict(vars(g._get_current_object())) if has_app_context() else {} + user = globals_snapshot.pop("user", None) if has_request_context() else None + globals_snapshot.pop("user", None) + # An expired ORM user's .id can issue metadata I/O on the transport loop. + # The identity key is available without loading any attributes. + user_state = sa_inspect(user, raiseerr=False) + user_id = ( + user_state.identity[0] + if isinstance(user_state, InstanceState) and user_state.identity + else getattr(user, "id", None) + ) + guest = user if getattr(user, "is_guest_user", False) else None + context = copy_context() + + def execute() -> Any: + """Own context teardown even when the caller stops waiting.""" + from contextlib import nullcontext + + from superset import db, security_manager + from superset.sql.execution.cancellation import cursor_scope + + _active_call.set(call) + cursor_scope.set(warehouse_cursor) + with _worker_context(app): + vars(g._get_current_object()).update(globals_snapshot) + with request_copy if request_copy is not None else nullcontext(): + if user_id is not None: + g.user = db.session.get(security_manager.user_model, user_id) + elif guest is not None: + g.user = guest + call.check() + result = asyncio.run(fn(*args, **worker_kwargs)) + call.check() + return result + + future = pool.submit(lambda: context.run(execute), call.finished) + 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 + finally: + from superset.mcp_service.auth import _mcp_user_id_var + + _mcp_user_id_var.set(call.user_id) + + +class QueryCancellation: + """Bridge engine cancellation to a separate, freshly scoped metadata session.""" + + def __init__( + self, + call: WorkerCall, + database: Database, + cursor: Any, + catalog: str | None, + schema: str | None, + ) -> None: + self.call = call + self.database = database + self.database_id = database.id + self.spec = database.db_engine_spec + self.cursor = cursor + self.catalog = catalog + self.schema = schema + self.cancel_id: str | None = None + self.context = copy_context() + + def capture(self) -> None: + """Capture an engine handle without making unsupported drivers fail.""" + from superset.tasks.query_cancel import capture_cancel_query_id + + try: + self.cancel_id = capture_cancel_query_id(self.database, self.cursor) + except Exception: + logger.warning( + "MCP call %s: cancel-id capture failed", + self.call.call_id, + exc_info=True, + ) + + def register(self) -> None: + """Publish a usable handle, including one obtained after the deadline.""" + with self.call.lock: + if self.spec.has_implicit_cancel(): Review Comment: A correction to my earlier answer: this is now implemented, in 2c1e351025. - `TrinoEngineSpec.get_cancel_query_id` returns `cursor.query_id`, which is `None` until execution starts. SQL Lab (`sql_lab.py`), the Celery SQL task and the chart-data task seam all ask before executing, so they record nothing new, and SQL Lab's `handle_cursor` / `QUERY_CANCEL_KEY` path is unchanged. - A new spec flag, `has_query_id_during_execute`, is True only for `TrinoEngineSpec`. With it, MCP registers cancellation without a handle. On abandonment the cancellation thread reads the queryId from the still-blocked cursor and calls the existing `cancel_chart_query` → `TrinoEngineSpec.cancel_query` (`system.runtime.kill_query`). - `test_abandoned_trino_query_is_killed_by_its_query_id` runs a real trino 0.339.0 DB-API cursor through `TrinoEngineSpec.execute` against a fake coordinator that keeps the query RUNNING. After the 0.3s deadline, the kill receives the coordinator's queryId on a thread other than the worker's. The worker's `execute` then ends with `ADMINISTRATIVELY_KILLED` for that queryId. Without the change the test fails because the worker stays blocked. - Limits: the kill call is mocked in that test, so I have not run it against a real Trino. BigQuery, Databricks native, Athena and SQLite specs still provide no cancel id; for them abandonment still logs `dispatched=False`, and the call keeps its slot until the driver returns. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
