eschutho commented on code in PR #42900:
URL: https://github.com/apache/superset/pull/42900#discussion_r4449115400
##########
superset/db_engine_specs/databricks.py:
##########
@@ -931,6 +950,254 @@ def adjust_engine_params(
return uri, connect_args
+ @staticmethod
+ def _uses_sea(cursor: Any) -> bool:
+ """Whether ``cursor`` runs over the Statement Execution API (SEA)."""
+ session = getattr(getattr(cursor, "connection", None), "session", None)
+ return bool(getattr(session, "use_sea", False))
+
+ @classmethod
+ def execute_with_cursor(cls, cursor: Any, sql: str, query: Query) -> None:
+ """
+ Run a statement so a stop can cancel it while it is still executing.
+
+ On SEA connections the statement is submitted with ``execute_async``,
+ which publishes its statement id before returning. The id is recorded
+ as the query's cancel id, and the statement is polled with
+ ``is_query_pending`` until it finishes or a stop is requested, in which
+ case the live cursor cancels it. Once the statement is no longer
+ running its id is retired, so a stop between statements of a
+ multi-statement script never targets an already finished statement.
+
+ Thrift connections keep the synchronous execution path, as their
+ statements can't be cancelled from the stop request's connection.
+ """
+ if not cls._uses_sea(cursor):
+ super().execute_with_cursor(cursor, sql, query)
+ return
+
+ logger.debug("Query %d: Running query: %s", query.id, sql)
+ if not query.extra.get(cls.sea_async_execution_key):
+ query.set_extra_json_key(cls.sea_async_execution_key, True)
+ db.session.commit() # pylint: disable=consider-using-transaction
+
+ if cls.arraysize:
+ cursor.arraysize = cls.arraysize
+ try:
+ cursor.execute_async(sql)
+ except Exception as ex:
+ database = query.database
+ if (
+ not is_oauth2_retry_active()
+ and database.is_oauth2_enabled()
+ and cls.needs_oauth2(ex)
+ ):
+ cls.start_oauth2_dance(database)
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ try:
+ if (cancel_query_id := cls.get_cancel_query_id(cursor, query)) is
not None:
+ query.set_extra_json_key(QUERY_CANCEL_KEY, cancel_query_id)
+ db.session.commit() # pylint:
disable=consider-using-transaction
+ stopped = cls._poll_statement(cursor, query)
+ finally:
+ cls._retire_cancel_query_id(query)
+
+ if stopped:
+ return
+
+ try:
+ # Loads the finished statement's result set onto the cursor, so it
+ # can be fetched exactly as after a synchronous ``execute``.
+ cursor.get_async_execution_result()
+ except Exception as ex:
+ try:
+ cls._raise_statement_error(cursor)
+ except Exception as statement_error: # pylint:
disable=broad-except
+ raise cls.get_dbapi_mapped_exception(
+ statement_error
+ ) from statement_error
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ @staticmethod
+ def _raise_statement_error(cursor: Any) -> None:
+ """
+ Raise the warehouse's error for a statement that didn't succeed.
+
+ ``get_async_execution_result`` reports any unsuccessful statement with
+ a generic "Operation status FAILED" error. The synchronous ``execute``
+ instead raises the server's error code and message, through the SEA
+ backend's state check. The driver exposes no public accessor for that
+ detail, so the same check is applied to the statement's latest status.
+ Returns without raising if the detail can't be retrieved, leaving the
+ caller to raise the driver's generic error.
+ """
+ backend = getattr(cursor, "backend", None)
+ poll_query = getattr(backend, "_poll_query", None)
+ check_state = getattr(
+ backend, "_check_command_not_in_failed_or_closed_state", None
+ )
+ command_id = getattr(cursor, "active_command_id", None)
+ if poll_query is None or check_state is None or command_id is None:
+ return
+ try:
+ status = poll_query(command_id).status
+ except SoftTimeLimitExceeded:
+ raise
+ except Exception: # pylint: disable=broad-except
+ return
+ check_state(status, command_id)
+
+ @classmethod
+ def _poll_statement(cls, cursor: Any, query: Query) -> bool:
+ """
+ Wait for the cursor's statement to finish, cancelling it on a stop.
+
+ A failed status request doesn't mean the statement finished, so it's
+ treated as still running and polled again; the statement stays
+ cancellable meanwhile. After ``sea_max_poll_failures`` consecutive
+ failures the statement is cancelled and the error is raised.
+
+ :return: True if the query was stopped, False if the statement finished
+ """
+ # Short statements finish within moments, so polling starts with a
+ # short interval that backs off to the configured one.
+ sleep_interval = 0.1
+ max_interval = app.config["DB_POLL_INTERVAL_SECONDS"].get(cls.engine,
1)
+ failures = 0
+ try:
+ while True:
+ try:
+ pending = cursor.is_query_pending()
+ failures = 0
+ except (SoftTimeLimitExceeded, SupersetTimeoutException):
+ raise
+ except Exception as ex: # pylint: disable=broad-except
+ failures += 1
+ if failures >= cls.sea_max_poll_failures:
+ cls._cancel_statement(cursor, query.id)
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+ logger.warning(
+ "Query %s: polling statement state failed, retrying",
+ query.id,
+ )
+ pending = True
+ db.session.refresh(query)
Review Comment:
Good catch, thanks @sadpandajoe. Fixed in 9b32f4f20c.
The loop now calls `db.session.commit()` before each
`db.session.refresh(query)`, the same as `ImpalaEngineSpec.handle_cursor`. Each
poll starts a fresh read transaction, so a Stop committed after polling begins
is seen and the live-cursor `cancel()` fires. `_retire_cancel_query_id` had the
same issue: it reloaded from the old snapshot before rewriting `extra`, which
could also overwrite an early-cancel flag committed in the meantime. It now
commits before reloading too.
Test:
`test_execute_with_cursor_sea_sees_stop_committed_after_polling_starts` uses a
session that serves reloads from the transaction's first snapshot, as
REPEATABLE READ does, and commits a Stop from "another request" mid-poll. It
asserts the live statement gets cancelled. It fails if the commit before
refresh is removed.
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -1646,3 +1646,556 @@ def
test_databricks_hive_spec_shares_extended_aggregations() -> None:
for aggregate in ("MEDIAN", "STDDEV_SAMP", "VAR_SAMP"):
assert
DatabricksHiveEngineSpec.get_extended_aggregation_func(aggregate)
+
+
+def test_has_query_id_before_execute_is_false() -> None:
+ """
+ The cancel id can only be read off the cursor after a statement has been
+ executed (``cursor.active_command_id`` is unset before that), so the id
+ must be captured post-execute rather than up front.
+ """
+ assert DatabricksPythonConnectorEngineSpec.has_query_id_before_execute is
False
+
+
+def test_get_cancel_query_id_sea(mocker: MockerFixture) -> None:
+ """
+ On a SEA (Statement Execution API) connection, the cancel id is the
+ statement id exposed by the driver's ``CommandId``.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id.to_sea_statement_id.return_value = "01ecc35f-abcd"
+ cursor.connection.session.use_sea = True
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
+ == "01ecc35f-abcd"
+ )
+
+
+def test_get_cancel_query_id_no_active_command(mocker: MockerFixture) -> None:
+ """
+ Before a statement has been executed on the cursor, there's nothing to
+ build a cancel id from.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id = None
+ cursor.connection.session.use_sea = True
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
is None
+ )
+
+
+def test_get_cancel_query_id_thrift_backend_returns_none(
+ mocker: MockerFixture,
+) -> None:
+ """
+ When the connection was not established with ``use_sea=True`` (i.e. it's
+ using the default Thrift backend), no cancel id is captured -- cancelling
+ a Thrift command requires a secret that's never exposed through any
+ public driver API, so we deliberately don't attempt it.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id.to_sea_statement_id.return_value = "01ecc35f-abcd"
+ cursor.connection.session.use_sea = False
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
is None
+ )
+
+
[email protected]
+def databricks_command_id(mocker: MockerFixture) -> Any:
+ """
+ Stand-in for ``databricks.sql.backend.types.CommandId``.
+
+ The optional ``databricks-sql-connector`` driver isn't installed in the
+ unit-test environment, so the module chain is injected into ``sys.modules``
+ to let ``cancel_query``'s inline import resolve. Yields the ``CommandId``
+ mock so tests can assert on ``from_sea_statement_id``.
+ """
+ command_id_cls = mocker.MagicMock()
+ mocker.patch.dict(
+ "sys.modules",
+ {
+ "databricks": mocker.MagicMock(),
+ "databricks.sql": mocker.MagicMock(),
+ "databricks.sql.backend": mocker.MagicMock(),
+ "databricks.sql.backend.types":
mocker.MagicMock(CommandId=command_id_cls),
+ },
+ )
+ return command_id_cls
+
+
+def test_cancel_query_sea_success(
+ mocker: MockerFixture, databricks_command_id: Any
+) -> None:
+ """
+ ``cancel_query`` sets the reconstructed SEA ``CommandId`` on the fresh
+ cursor and delegates to the driver's own ``cursor.cancel()``.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.cancel_query(cursor, query,
"01ecc35f-abcd")
+ is True
+ )
+
databricks_command_id.from_sea_statement_id.assert_called_once_with("01ecc35f-abcd")
+ assert (
+ cursor.active_command_id
+ == databricks_command_id.from_sea_statement_id.return_value
+ )
+ cursor.cancel.assert_called_once_with()
+
+
+def test_cancel_query_invalid_id_returns_false(mocker: MockerFixture) -> None:
+ """
+ A malformed cancel id is rejected before touching the cursor at all.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.cancel_query(
+ cursor, query, "'; DROP TABLE foo; --"
+ )
+ is False
+ )
+ cursor.cancel.assert_not_called()
+
+
+def test_cancel_query_propagates_errors(
+ mocker: MockerFixture, databricks_command_id: Any
+) -> None:
+ """
+ If the cancel attempt itself fails (e.g. the fresh cursor/connection
+ errors), the error must surface rather than being reported as a failed-
+ but-handled cancel.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.cancel.side_effect = RuntimeError("connection reset")
+
+ query = Query()
+ with pytest.raises(RuntimeError):
+ DatabricksPythonConnectorEngineSpec.cancel_query(cursor, query,
"01ecc35f-abcd")
+
+
+class FakeSeaCursor:
+ """
+ Minimal SEA cursor: each ``execute_async`` publishes the next statement id,
+ and the statement reports as pending for a fixed number of polls.
+ """
+
+ def __init__(self, statement_ids: list[str], pending_polls: int) -> None:
+ self.connection = MagicMock()
+ self.connection.session.use_sea = True
+ self.backend = MagicMock()
+ self.active_command_id: Any = None
+ self.statement_ids = iter(statement_ids)
+ self.pending_polls = pending_polls
+ self.statement_id: str | None = None
+ self.remaining_polls = 0
+ self.cancelled: list[str | None] = []
+ self.results_loaded: list[str | None] = []
+ # Exceptions raised by successive status polls before they succeed.
+ self.poll_errors: list[BaseException] = []
+
+ def execute_async(self, sql: str) -> None:
+ self.statement_id = next(self.statement_ids)
+ self.active_command_id = MagicMock()
+ self.active_command_id.to_sea_statement_id.return_value =
self.statement_id
+ self.remaining_polls = self.pending_polls
+
+ def is_query_pending(self) -> bool:
+ if self.poll_errors:
+ raise self.poll_errors.pop(0)
+ if self.remaining_polls > 0:
+ self.remaining_polls -= 1
+ return True
+ return False
+
+ def cancel(self) -> None:
+ self.cancelled.append(self.statement_id)
+ self.remaining_polls = 0
+
+ def get_async_execution_result(self) -> None:
+ self.results_loaded.append(self.statement_id)
+
+
[email protected]
+def sea_session(mocker: MockerFixture) -> Any:
+ """Stub the metadata session and polling clock used by the SEA loop."""
+ mocker.patch("superset.db_engine_specs.databricks.time.sleep")
+ mocker.patch(
+ "superset.db_engine_specs.databricks.app",
+ config={"DB_POLL_INTERVAL_SECONDS": {}},
+ )
+ return mocker.patch("superset.db_engine_specs.databricks.db")
+
+
+def test_execute_with_cursor_sea_finished_statement(sea_session: Any) -> None:
+ """
+ A statement that finishes on its own has its results loaded onto the
+ cursor, and its id is retired once it's no longer running.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=2)
+ query = Query(status=QueryStatus.RUNNING)
+ published: list[Any] = []
+ sea_session.session.refresh.side_effect = lambda q: published.append(
+ q.extra.get(QUERY_CANCEL_KEY)
+ )
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ assert "stmt-1" in published
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == []
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_cancels_in_flight_statement(
+ sea_session: Any,
+) -> None:
+ """
+ A stop that lands while the statement is still running on the warehouse
+ cancels it through the live cursor, and no results are loaded.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ query = Query(status=QueryStatus.RUNNING)
+ polls = 0
+
+ def refresh(q: Query) -> None:
+ nonlocal polls
+ polls += 1
+ if polls == 3:
+ # The stop is persisted while the statement is mid-flight, and its
+ # published id is the one a stop request would cancel.
+ assert q.extra.get(QUERY_CANCEL_KEY) == "stmt-1"
+ q.status = QueryStatus.STOPPED
+
+ sea_session.session.refresh.side_effect = refresh
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ assert cursor.cancelled == ["stmt-1"]
+ assert cursor.remaining_polls == 0
+ assert cursor.results_loaded == []
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_multi_statement_cancels_current(
+ sea_session: Any,
+) -> None:
+ """
+ On a multi-statement script, a stop cancels whichever statement is running:
+ the finished statement's id is retired and never targeted.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1", "stmt-2"], pending_polls=5)
+ query = Query(status=QueryStatus.RUNNING)
+ published: list[Any] = []
+ stop_on_statement = "stmt-2"
+
+ def refresh(q: Query) -> None:
+ published.append(q.extra.get(QUERY_CANCEL_KEY))
+ if q.extra.get(QUERY_CANCEL_KEY) == stop_on_statement:
+ q.status = QueryStatus.STOPPED
+
+ sea_session.session.refresh.side_effect = refresh
+
+ spec = DatabricksPythonConnectorEngineSpec
+ spec.execute_with_cursor(cursor, "SELECT 1", query)
+ # Between the statements no id is published, so nothing stale is cancelled.
+ assert QUERY_CANCEL_KEY not in query.extra
+ spec.execute_with_cursor(cursor, "SELECT 2", query)
+
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == ["stmt-2"]
+ assert "stmt-1" in published
+ assert "stmt-2" in published
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_stop_between_statements(sea_session: Any) ->
None:
+ """
+ A stop that arrives between statements finds no published id; it is routed
+ to the polling loop, which cancels the next statement as soon as it polls.
+ """
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1", "stmt-2"], pending_polls=5)
+ query = Query(status=QueryStatus.RUNNING)
+
+ spec = DatabricksPythonConnectorEngineSpec
+ spec.execute_with_cursor(cursor, "SELECT 1", query)
+
+ spec.prepare_cancel_query(query)
+ assert query.extra.get(QUERY_EARLY_CANCEL_KEY) is True
+
+ spec.execute_with_cursor(cursor, "SELECT 2", query)
+
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == ["stmt-2"]
+
+
+def test_prepare_cancel_query_with_published_id_is_noop() -> None:
+ """While a statement id is published, the stop cancels it by id."""
+ from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.set_extra_json_key(
+ DatabricksPythonConnectorEngineSpec.sea_async_execution_key, True
+ )
+ query.set_extra_json_key(QUERY_CANCEL_KEY, "stmt-1")
+
+ DatabricksPythonConnectorEngineSpec.prepare_cancel_query(query)
+
+ assert QUERY_EARLY_CANCEL_KEY not in query.extra
+
+
+def test_prepare_cancel_query_thrift_is_noop() -> None:
+ """
+ Thrift statements never run through the polling loop, so a stop is not
+ reported as handled by staging early-cancel intent.
+ """
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ DatabricksPythonConnectorEngineSpec.prepare_cancel_query(query)
+
+ assert QUERY_EARLY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_thrift_uses_synchronous_path(
+ mocker: MockerFixture,
+) -> None:
+ """Thrift connections keep the base synchronous execution."""
+ from superset.models.sql_lab import Query
+
+ base = mocker.patch(
+ "superset.db_engine_specs.databricks."
+ "DatabricksDynamicBaseEngineSpec.execute_with_cursor"
+ )
+ cursor = mocker.MagicMock()
+ cursor.connection.session.use_sea = False
+ query = Query()
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ base.assert_called_once_with(cursor, "SELECT 1", query)
+ cursor.execute_async.assert_not_called()
+
+
+def test_execute_with_cursor_sea_soft_time_limit_cancels(sea_session: Any) ->
None:
+ """Hitting the task's soft time limit cancels the running statement."""
+ from celery.exceptions import SoftTimeLimitExceeded
+
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ query = Query(status=QueryStatus.RUNNING)
+ polls = 0
+
+ def refresh(q: Query) -> None:
+ nonlocal polls
+ polls += 1
+ if polls == 2:
+ raise SoftTimeLimitExceeded()
+
+ sea_session.session.refresh.side_effect = refresh
+
+ with pytest.raises(SoftTimeLimitExceeded):
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(
+ cursor, "SELECT 1", query
+ )
+
+ assert cursor.cancelled == ["stmt-1"]
+
+
+def test_execute_with_cursor_sea_transient_poll_failure_stays_cancellable(
+ sea_session: Any,
+) -> None:
+ """
+ A failed status request doesn't count as the statement finishing: its id
+ stays published and a stop that follows still cancels the live statement.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ cursor.poll_errors = [ConnectionResetError("connection reset by peer")]
+ query = Query(status=QueryStatus.RUNNING)
+ published: list[Any] = []
+
+ def refresh(q: Query) -> None:
+ published.append(q.extra.get(QUERY_CANCEL_KEY))
+ if len(published) == 3:
+ q.status = QueryStatus.STOPPED
+
+ sea_session.session.refresh.side_effect = refresh
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ # The poll after the failure still saw the statement's published id.
+ assert published[:3] == ["stmt-1", "stmt-1", "stmt-1"]
+ assert cursor.cancelled == ["stmt-1"]
+ assert cursor.results_loaded == []
+
+
+def test_execute_with_cursor_sea_transient_poll_failure_then_finishes(
+ sea_session: Any,
+) -> None:
+ """Polling resumes after a failed status request and loads the results."""
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=2)
+ cursor.poll_errors = [RuntimeError("503 Service Unavailable")]
+ query = Query(status=QueryStatus.RUNNING)
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ assert cursor.cancelled == []
+ assert cursor.results_loaded == ["stmt-1"]
+
+
+def test_execute_with_cursor_sea_persistent_poll_failure_cancels(
+ sea_session: Any,
+) -> None:
+ """
+ When status requests keep failing, the statement is cancelled rather than
+ left running unmonitored, and the error is raised.
+ """
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ spec = DatabricksPythonConnectorEngineSpec
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ cursor.poll_errors = [
Review Comment:
Agreed, the test couldn't tell those cases apart. Fixed in 9b32f4f20c.
- `test_execute_with_cursor_sea_poll_failure_count_resets_on_success`
scripts 9 failures, then a successful poll where the statement is still
pending, then 9 more failures. It asserts nothing is cancelled, the results
load, and the poll count is 20 (9 + 1 + 9 + the poll that sees it finished).
- `test_execute_with_cursor_sea_persistent_poll_failure_cancels` now
supplies 15 consecutive failures and asserts the statement is cancelled on
exactly the 10th poll (`cursor.polls == sea_max_poll_failures`). So giving up
early or late both fail.
I checked both against mutated code. Removing the counter reset breaks the
first test, and capping at 2 failures breaks the second.
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -1646,3 +1646,556 @@ def
test_databricks_hive_spec_shares_extended_aggregations() -> None:
for aggregate in ("MEDIAN", "STDDEV_SAMP", "VAR_SAMP"):
assert
DatabricksHiveEngineSpec.get_extended_aggregation_func(aggregate)
+
+
+def test_has_query_id_before_execute_is_false() -> None:
+ """
+ The cancel id can only be read off the cursor after a statement has been
+ executed (``cursor.active_command_id`` is unset before that), so the id
+ must be captured post-execute rather than up front.
+ """
+ assert DatabricksPythonConnectorEngineSpec.has_query_id_before_execute is
False
+
+
+def test_get_cancel_query_id_sea(mocker: MockerFixture) -> None:
+ """
+ On a SEA (Statement Execution API) connection, the cancel id is the
+ statement id exposed by the driver's ``CommandId``.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id.to_sea_statement_id.return_value = "01ecc35f-abcd"
+ cursor.connection.session.use_sea = True
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
+ == "01ecc35f-abcd"
+ )
+
+
+def test_get_cancel_query_id_no_active_command(mocker: MockerFixture) -> None:
+ """
+ Before a statement has been executed on the cursor, there's nothing to
+ build a cancel id from.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id = None
+ cursor.connection.session.use_sea = True
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
is None
+ )
+
+
+def test_get_cancel_query_id_thrift_backend_returns_none(
+ mocker: MockerFixture,
+) -> None:
+ """
+ When the connection was not established with ``use_sea=True`` (i.e. it's
+ using the default Thrift backend), no cancel id is captured -- cancelling
+ a Thrift command requires a secret that's never exposed through any
+ public driver API, so we deliberately don't attempt it.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.active_command_id.to_sea_statement_id.return_value = "01ecc35f-abcd"
+ cursor.connection.session.use_sea = False
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.get_cancel_query_id(cursor, query)
is None
+ )
+
+
[email protected]
+def databricks_command_id(mocker: MockerFixture) -> Any:
+ """
+ Stand-in for ``databricks.sql.backend.types.CommandId``.
+
+ The optional ``databricks-sql-connector`` driver isn't installed in the
+ unit-test environment, so the module chain is injected into ``sys.modules``
+ to let ``cancel_query``'s inline import resolve. Yields the ``CommandId``
+ mock so tests can assert on ``from_sea_statement_id``.
+ """
+ command_id_cls = mocker.MagicMock()
+ mocker.patch.dict(
+ "sys.modules",
+ {
+ "databricks": mocker.MagicMock(),
+ "databricks.sql": mocker.MagicMock(),
+ "databricks.sql.backend": mocker.MagicMock(),
+ "databricks.sql.backend.types":
mocker.MagicMock(CommandId=command_id_cls),
+ },
+ )
+ return command_id_cls
+
+
+def test_cancel_query_sea_success(
+ mocker: MockerFixture, databricks_command_id: Any
+) -> None:
+ """
+ ``cancel_query`` sets the reconstructed SEA ``CommandId`` on the fresh
+ cursor and delegates to the driver's own ``cursor.cancel()``.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.cancel_query(cursor, query,
"01ecc35f-abcd")
+ is True
+ )
+
databricks_command_id.from_sea_statement_id.assert_called_once_with("01ecc35f-abcd")
+ assert (
+ cursor.active_command_id
+ == databricks_command_id.from_sea_statement_id.return_value
+ )
+ cursor.cancel.assert_called_once_with()
+
+
+def test_cancel_query_invalid_id_returns_false(mocker: MockerFixture) -> None:
+ """
+ A malformed cancel id is rejected before touching the cursor at all.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ query = Query()
+ assert (
+ DatabricksPythonConnectorEngineSpec.cancel_query(
+ cursor, query, "'; DROP TABLE foo; --"
+ )
+ is False
+ )
+ cursor.cancel.assert_not_called()
+
+
+def test_cancel_query_propagates_errors(
+ mocker: MockerFixture, databricks_command_id: Any
+) -> None:
+ """
+ If the cancel attempt itself fails (e.g. the fresh cursor/connection
+ errors), the error must surface rather than being reported as a failed-
+ but-handled cancel.
+ """
+ from superset.models.sql_lab import Query
+
+ cursor = mocker.MagicMock()
+ cursor.cancel.side_effect = RuntimeError("connection reset")
+
+ query = Query()
+ with pytest.raises(RuntimeError):
+ DatabricksPythonConnectorEngineSpec.cancel_query(cursor, query,
"01ecc35f-abcd")
+
+
+class FakeSeaCursor:
+ """
+ Minimal SEA cursor: each ``execute_async`` publishes the next statement id,
+ and the statement reports as pending for a fixed number of polls.
+ """
+
+ def __init__(self, statement_ids: list[str], pending_polls: int) -> None:
+ self.connection = MagicMock()
+ self.connection.session.use_sea = True
+ self.backend = MagicMock()
+ self.active_command_id: Any = None
+ self.statement_ids = iter(statement_ids)
+ self.pending_polls = pending_polls
+ self.statement_id: str | None = None
+ self.remaining_polls = 0
+ self.cancelled: list[str | None] = []
+ self.results_loaded: list[str | None] = []
+ # Exceptions raised by successive status polls before they succeed.
+ self.poll_errors: list[BaseException] = []
+
+ def execute_async(self, sql: str) -> None:
+ self.statement_id = next(self.statement_ids)
+ self.active_command_id = MagicMock()
+ self.active_command_id.to_sea_statement_id.return_value =
self.statement_id
+ self.remaining_polls = self.pending_polls
+
+ def is_query_pending(self) -> bool:
+ if self.poll_errors:
+ raise self.poll_errors.pop(0)
+ if self.remaining_polls > 0:
+ self.remaining_polls -= 1
+ return True
+ return False
+
+ def cancel(self) -> None:
+ self.cancelled.append(self.statement_id)
+ self.remaining_polls = 0
+
+ def get_async_execution_result(self) -> None:
+ self.results_loaded.append(self.statement_id)
+
+
[email protected]
+def sea_session(mocker: MockerFixture) -> Any:
+ """Stub the metadata session and polling clock used by the SEA loop."""
+ mocker.patch("superset.db_engine_specs.databricks.time.sleep")
+ mocker.patch(
+ "superset.db_engine_specs.databricks.app",
+ config={"DB_POLL_INTERVAL_SECONDS": {}},
+ )
+ return mocker.patch("superset.db_engine_specs.databricks.db")
+
+
+def test_execute_with_cursor_sea_finished_statement(sea_session: Any) -> None:
+ """
+ A statement that finishes on its own has its results loaded onto the
+ cursor, and its id is retired once it's no longer running.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=2)
+ query = Query(status=QueryStatus.RUNNING)
+ published: list[Any] = []
+ sea_session.session.refresh.side_effect = lambda q: published.append(
+ q.extra.get(QUERY_CANCEL_KEY)
+ )
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ assert "stmt-1" in published
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == []
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_cancels_in_flight_statement(
+ sea_session: Any,
+) -> None:
+ """
+ A stop that lands while the statement is still running on the warehouse
+ cancels it through the live cursor, and no results are loaded.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ query = Query(status=QueryStatus.RUNNING)
+ polls = 0
+
+ def refresh(q: Query) -> None:
+ nonlocal polls
+ polls += 1
+ if polls == 3:
+ # The stop is persisted while the statement is mid-flight, and its
+ # published id is the one a stop request would cancel.
+ assert q.extra.get(QUERY_CANCEL_KEY) == "stmt-1"
+ q.status = QueryStatus.STOPPED
+
+ sea_session.session.refresh.side_effect = refresh
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ assert cursor.cancelled == ["stmt-1"]
+ assert cursor.remaining_polls == 0
+ assert cursor.results_loaded == []
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_multi_statement_cancels_current(
+ sea_session: Any,
+) -> None:
+ """
+ On a multi-statement script, a stop cancels whichever statement is running:
+ the finished statement's id is retired and never targeted.
+ """
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1", "stmt-2"], pending_polls=5)
+ query = Query(status=QueryStatus.RUNNING)
+ published: list[Any] = []
+ stop_on_statement = "stmt-2"
+
+ def refresh(q: Query) -> None:
+ published.append(q.extra.get(QUERY_CANCEL_KEY))
+ if q.extra.get(QUERY_CANCEL_KEY) == stop_on_statement:
+ q.status = QueryStatus.STOPPED
+
+ sea_session.session.refresh.side_effect = refresh
+
+ spec = DatabricksPythonConnectorEngineSpec
+ spec.execute_with_cursor(cursor, "SELECT 1", query)
+ # Between the statements no id is published, so nothing stale is cancelled.
+ assert QUERY_CANCEL_KEY not in query.extra
+ spec.execute_with_cursor(cursor, "SELECT 2", query)
+
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == ["stmt-2"]
+ assert "stmt-1" in published
+ assert "stmt-2" in published
+ assert QUERY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_sea_stop_between_statements(sea_session: Any) ->
None:
+ """
+ A stop that arrives between statements finds no published id; it is routed
+ to the polling loop, which cancels the next statement as soon as it polls.
+ """
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1", "stmt-2"], pending_polls=5)
+ query = Query(status=QueryStatus.RUNNING)
+
+ spec = DatabricksPythonConnectorEngineSpec
+ spec.execute_with_cursor(cursor, "SELECT 1", query)
+
+ spec.prepare_cancel_query(query)
+ assert query.extra.get(QUERY_EARLY_CANCEL_KEY) is True
+
+ spec.execute_with_cursor(cursor, "SELECT 2", query)
+
+ assert cursor.results_loaded == ["stmt-1"]
+ assert cursor.cancelled == ["stmt-2"]
+
+
+def test_prepare_cancel_query_with_published_id_is_noop() -> None:
+ """While a statement id is published, the stop cancels it by id."""
+ from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ query.set_extra_json_key(
+ DatabricksPythonConnectorEngineSpec.sea_async_execution_key, True
+ )
+ query.set_extra_json_key(QUERY_CANCEL_KEY, "stmt-1")
+
+ DatabricksPythonConnectorEngineSpec.prepare_cancel_query(query)
+
+ assert QUERY_EARLY_CANCEL_KEY not in query.extra
+
+
+def test_prepare_cancel_query_thrift_is_noop() -> None:
+ """
+ Thrift statements never run through the polling loop, so a stop is not
+ reported as handled by staging early-cancel intent.
+ """
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+
+ query = Query()
+ DatabricksPythonConnectorEngineSpec.prepare_cancel_query(query)
+
+ assert QUERY_EARLY_CANCEL_KEY not in query.extra
+
+
+def test_execute_with_cursor_thrift_uses_synchronous_path(
+ mocker: MockerFixture,
+) -> None:
+ """Thrift connections keep the base synchronous execution."""
+ from superset.models.sql_lab import Query
+
+ base = mocker.patch(
+ "superset.db_engine_specs.databricks."
+ "DatabricksDynamicBaseEngineSpec.execute_with_cursor"
+ )
+ cursor = mocker.MagicMock()
+ cursor.connection.session.use_sea = False
+ query = Query()
+
+ DatabricksPythonConnectorEngineSpec.execute_with_cursor(cursor, "SELECT
1", query)
+
+ base.assert_called_once_with(cursor, "SELECT 1", query)
+ cursor.execute_async.assert_not_called()
+
+
+def test_execute_with_cursor_sea_soft_time_limit_cancels(sea_session: Any) ->
None:
+ """Hitting the task's soft time limit cancels the running statement."""
+ from celery.exceptions import SoftTimeLimitExceeded
+
+ from superset.models.sql_lab import Query
+ from superset.utils.core import QueryStatus
+
+ cursor = FakeSeaCursor(["stmt-1"], pending_polls=100)
+ query = Query(status=QueryStatus.RUNNING)
+ polls = 0
+
+ def refresh(q: Query) -> None:
+ nonlocal polls
+ polls += 1
+ if polls == 2:
+ raise SoftTimeLimitExceeded()
Review Comment:
Added in 9b32f4f20c:
`test_execute_with_cursor_sea_soft_time_limit_during_poll_cancels` puts a
`SoftTimeLimitExceeded` in `FakeSeaCursor.poll_errors`. It asserts the
exception propagates immediately, after exactly one poll and with exactly one
`cancelled` entry. It fails if `SoftTimeLimitExceeded` is dropped from the
inner `except ...: raise` clause.
--
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]
