aminghadersohi commented on code in PR #44581:
URL: https://github.com/apache/superset/pull/44581#discussion_r4090574075
##########
superset/mcp_service/auth.py:
##########
@@ -1259,7 +1277,27 @@ async def async_wrapper(*args: Any, **kwargs: Any) ->
Any:
_cleanup_session_on_error()
raise
- wrapper = async_wrapper
+ @functools.wraps(tool_func)
+ async def worker_wrapper(*args: Any, **kwargs: Any) -> Any:
+ from superset.mcp_service.worker import run_in_worker
+
+ bound = _tool_sig.bind_partial(*args, **kwargs)
+ request = bound.arguments.get("request")
+ seconds = getattr(request, "timeout", None)
+ if seconds is None:
+ if has_app_context():
+ seconds = current_app.config.get("SQLLAB_TIMEOUT", 30)
+ else:
+ from superset.mcp_service.flask_singleton import
get_flask_app
+
+ seconds = get_flask_app().config.get("SQLLAB_TIMEOUT", 30)
+ # Bind ctx once on the transport loop, including positional
callers.
+ bound.arguments.update(_inject_ctx(dict(bound.arguments)))
+ return await run_in_worker(
+ async_wrapper, (), dict(bound.arguments), seconds
+ )
+
+ wrapper = worker_wrapper
Review Comment:
Fixed in 19bb267df3 (direct-call test updates in f6d2f2d6b7). There are
eleven registered synchronous tools; update_chart_preview reaches warehouse
execution through generate_preview_from_form_data -> ChartDataCommand.run.
FastMCP can offload sync calls itself, but that bypassed this PR’s bounded
workers, deadline, and cancellation context. Both protected function kinds now
use run_in_worker. The new sync chart/SQL Lab regressions failed before the fix
and pass afterward; they check deadline responses and concurrent request
progress. The complete MCP suite plus related modules passes (5,360 tests; four
existing skips), and pre-commit passes.
##########
superset/models/core.py:
##########
@@ -602,35 +602,38 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
sqlalchemy_uri=sqlalchemy_uri,
cacheable=not prequeries,
)
- if prequeries:
- # SQLAlchemy connect event: runs prequeries on every
new
- # DBAPI connection (e.g. SET search_path for
PostgreSQL).
- def run_prequeries(
- dbapi_connection: Any,
- connection_record: Any, # pylint:
disable=unused-argument
- ) -> None:
- cursor = dbapi_connection.cursor()
+ from superset.sql.execution.cancellation import
cancellable_engine
+
+ with cancellable_engine(self, engine, catalog, schema):
+ if prequeries:
+ # SQLAlchemy connect event: runs prequeries on
every new
+ # DBAPI connection (e.g. SET search_path for
PostgreSQL).
+ def run_prequeries(
+ dbapi_connection: Any,
+ connection_record: Any, # pylint:
disable=unused-argument
+ ) -> None:
+ cursor = dbapi_connection.cursor()
+ try:
+ for prequery in prequeries:
+ cursor.execute(prequery)
Review Comment:
Fixed in 19bb267df3. The raw connect-event prequery cursor now enters
cancellable_cursor before execution, checks the deadline before each statement,
and refreshes cancellation handles afterward. The regression uses a blocking
DBAPI prequery: it failed before the fix because cancellation was never
dispatched; it now verifies cancellation, cursor cleanup, and that the next
prequery never executes after timeout. The SQL suite also passes its 100%
coverage gate (1,975 tests).
--
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]
