rebenitez1802 commented on code in PR #45729:
URL: https://github.com/apache/superset/pull/45729#discussion_r4402097587
##########
superset/mcp_service/worker.py:
##########
@@ -791,7 +793,10 @@ def capture(self) -> None:
def register(self) -> None:
"""Publish a usable handle, including one obtained after the
deadline."""
with self.call.lock:
- if self.spec.has_implicit_cancel():
+ if self.spec.has_implicit_cancel() or (
+ self.spec.supports_cancel_on_live_cursor is True
+ and callable(getattr(self.cursor, "cancel", None))
+ ):
Review Comment:
🔴 **High: cross-thread cancel isn't limited to pyodbc**
The new docs promise this only for `pyodbc`, but this gate only checks for a
callable `cursor.cancel`. `MssqlEngineSpec` declares no `drivers`, so
`mssql+pytds` URLs resolve to it too. python-tds' `Cursor.cancel()` calls
`cancel_if_pending()`, which sends the cancel request and then reads the reply
from the same socket the query thread is already blocked reading
(`pytds/tds_session.py` `cancel_if_pending` / `process_cancel`). That socket
has no timeout by default. So when an MCP deadline fires, one of the two
threads can wait forever, and the call's worker slot isn't freed until a
restart. On master SQL Server had no MCP cancel at all, so this hang is new.
Suggested fix: let the engine spec decide from the driver. Add
`can_cancel_live_cursor(driver, cursor)` to `BaseEngineSpec` and override it in
`MssqlEngineSpec` to require `pyodbc`. The suggestions on `base.py`, `mssql.py`
and `test_worker.py` go with this one:
```suggestion
live = self.spec.can_cancel_live_cursor(self.database.driver,
self.cursor)
with self.call.lock:
if self.spec.has_implicit_cancel() or live is True:
```
##########
superset/db_engine_specs/base.py:
##########
@@ -816,6 +816,10 @@ def get_extended_aggregation_func(
# owner cancel a query whose caller has given up.
has_query_id_during_execute = False
+ # Execution owners may call cursor.cancel() from another thread when the
+ # driver exposes it. This does not enable SQL Lab's cross-process stop
path.
+ supports_cancel_on_live_cursor = False
Review Comment:
Part of the 🔴 fix on `QueryCancellation.register()`: the default keeps the
existing flag and callable check, and specs can narrow it by driver.
```suggestion
supports_cancel_on_live_cursor = False
@classmethod
def can_cancel_live_cursor( # pylint: disable=unused-argument
cls, driver: str, cursor: Any
) -> bool:
"""Whether ``cursor.cancel()`` is safe to call from another
thread."""
return cls.supports_cancel_on_live_cursor is True and callable(
getattr(cursor, "cancel", None)
)
```
##########
superset/db_engine_specs/mssql.py:
##########
@@ -96,6 +96,7 @@ class MssqlEngineSpec(BaseEngineSpec):
max_column_name_length = 128
allows_cte_in_subquery = False
supports_multivalues_insert = True
+ supports_cancel_on_live_cursor = True
Review Comment:
Part of the 🔴 fix on `QueryCancellation.register()`: only pyodbc's
`cancel()` (ODBC `SQLCancel`) is designed to be called from another thread.
python-tds' reads the reply from the socket the query thread is blocked on, and
pymssql's cursor has no `cancel()`. `AzureSynapseSpec` still opts out through
`super()`.
```suggestion
supports_cancel_on_live_cursor = True
@classmethod
def can_cancel_live_cursor(cls, driver: str, cursor: Any) -> bool:
"""Only pyodbc's ``cancel()`` (ODBC SQLCancel) is safe across
threads."""
return driver == "pyodbc" and super().can_cancel_live_cursor(driver,
cursor)
```
##########
tests/unit_tests/mcp_service/test_worker.py:
##########
@@ -418,6 +418,103 @@ def
test_query_cancellation_capture_resets_stale_id_on_failed_recapture() -> Non
assert call.cancel_query is None
[email protected]("native_cancel", [False, True])
+def test_sql_server_cancellation_uses_only_the_live_cursor(native_cancel:
bool) -> None:
+ """SQL Server cancellation requires a driver with a native cursor API."""
+ from superset.db_engine_specs.mssql import MssqlEngineSpec
+ from superset.mcp_service.worker import _active_call, warehouse_cursor
+
+ call = Mock(lock=threading.RLock(), expired=threading.Event())
+ database = Mock(id=42, db_engine_spec=MssqlEngineSpec)
+ cursor = Mock(spec=["execute", "cancel"] if native_cancel else ["execute"])
+ token = _active_call.set(call)
+ try:
+ with warehouse_cursor(database, cursor, None, None):
+ if native_cancel:
Review Comment:
Part of the 🔴 fix on `QueryCancellation.register()`: adds a `pytds` case.
That case fails on the current head because the callback gets registered.
```suggestion
@pytest.mark.parametrize(
("driver", "native_cancel", "registered"),
[("pyodbc", False, False), ("pyodbc", True, True), ("pytds", True,
False)],
)
def test_sql_server_cancellation_uses_only_the_live_cursor(
driver: str, native_cancel: bool, registered: bool
) -> None:
"""SQL Server cancellation requires pyodbc's native cursor API."""
from superset.db_engine_specs.mssql import MssqlEngineSpec
from superset.mcp_service.worker import _active_call, warehouse_cursor
call = Mock(lock=threading.RLock(), expired=threading.Event())
database = Mock(id=42, db_engine_spec=MssqlEngineSpec, driver=driver)
cursor = Mock(spec=["execute", "cancel"] if native_cancel else
["execute"])
token = _active_call.set(call)
try:
with warehouse_cursor(database, cursor, None, None):
if registered:
```
--
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]
