aminghadersohi commented on code in PR #44737:
URL: https://github.com/apache/superset/pull/44737#discussion_r4140726042
##########
superset/db_engine_specs/impala.py:
##########
@@ -105,14 +107,7 @@ def get_schema_names(cls, inspector: Inspector) ->
set[str]:
@classmethod
def has_implicit_cancel(cls) -> bool:
- """
- Return True if the live cursor handles the implicit cancelation of the
query,
- False otherwise.
-
- :return: Whether the live cursor implicitly cancels the query
- :see: handle_cursor
- """
-
+ """Keep HTTP cancellation independent of the live cursor polling
loop."""
Review Comment:
Fixed in 1d358309c4. `ImpalaEngineSpec.prepare_cancel_query()` now sets
`QUERY_EARLY_CANCEL_KEY` (and commits, as `trino.py` does) when no
`QUERY_CANCEL_KEY` has been published, so `cancel_query()` returns True and
`handle_cursor` cancels the live operation on its next poll.
`test_stop_before_cancel_id_is_published_uses_live_cursor` covers a dispatched
query with no cancel key through both `sql_lab.cancel_query` and
`SQLExecutor._cancel_query`: Stop succeeds, no HTTP call is made, and the
pending operation gets `cancel_operation()`. It fails without the change. On
the stale-ID case: `execute_with_cursor` re-captures the handle right after
`execute_async` (Impala has `has_query_id_before_execute = False`), so the
`USE` operation's ID is only stored in that same pre-execute window. If Stop's
HTTP cancel against it returns 200, the row is persisted STOPPED and the loop's
STOPPED check cancels the live operation.
##########
superset/db_engine_specs/impala.py:
##########
@@ -134,6 +129,7 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None:
query_id = query.id
unfinished_states = (
+ "PENDING_STATE",
Review Comment:
Good point. In 1d358309c4 the loop uses a short interval while the operation
is pending or initialized: 0.1s, doubling on each poll up to
`DB_POLL_INTERVAL_SECONDS`. A query that clears admission quickly is picked up
almost immediately, and one queued for longer doesn't hit the metadata DB every
100ms. RUNNING keeps the configured interval, as on master.
`test_pending_operation_polls_with_backoff` asserts the sleeps (0.1, 0.2, then
the configured 0.3 cap) and fails without the change.
##########
superset/db_engine_specs/impala.py:
##########
@@ -147,33 +143,37 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None:
# the query was still executed
# modified in stop_query in views / core.py is reflected here.
# stop query
- if query.extra.get(QUERY_EARLY_CANCEL_KEY):
- cursor.cancel_operation()
- cursor.close_operation()
- cursor.close()
+ if query.extra.get(QUERY_EARLY_CANCEL_KEY) or query.status in (
+ QueryStatus.STOPPED,
+ QueryStatus.TIMED_OUT,
Review Comment:
Right, the broad `except` swallowed it. Fixed in 1d358309c4: `handle_cursor`
catches `SoftTimeLimitExceeded` first, cancels and releases the operation, and
re-raises, so SQL Lab's handler marks the query TIMED_OUT. I dropped the
TIMED_OUT branch from the status check, since nothing can persist it while this
loop runs, and replaced the mock-status test with
`test_soft_time_limit_cancels_the_operation`. That test raises the exception
from the poll sleep in both PENDING and RUNNING and asserts the cancel plus the
re-raise. It fails without the change.
##########
superset/db_engine_specs/impala.py:
##########
@@ -105,14 +107,7 @@ def get_schema_names(cls, inspector: Inspector) ->
set[str]:
@classmethod
def has_implicit_cancel(cls) -> bool:
- """
- Return True if the live cursor handles the implicit cancelation of the
query,
- False otherwise.
-
- :return: Whether the live cursor implicitly cancels the query
- :see: handle_cursor
- """
-
+ """Keep HTTP cancellation independent of the live cursor polling
loop."""
return False
Review Comment:
The explicit-ID path does work once the operation is running: Impala sets
`has_query_id_before_execute = False`, so `BaseEngineSpec.execute_with_cursor`
calls `get_cancel_query_id` again right after `execute_async` and stores the
live handle before `handle_cursor` runs (`test_stop_with_cancel_id_uses_http`
covers that path). The real gap was a Stop landing before that handle is
published. 1d358309c4 closes it: `prepare_cancel_query()` sets the early-cancel
flag when no cancel key exists, so the stop succeeds and the live cursor
cancels the operation
(`test_stop_before_cancel_id_is_published_uses_live_cursor`).
##########
tests/unit_tests/db_engine_specs/test_impala_pending.py:
##########
@@ -0,0 +1,137 @@
+# 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.
+
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from flask import Flask
+
+from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
+from superset.db_engine_specs.impala import ImpalaEngineSpec
+from superset.sql.execution.executor import SQLExecutor
+from superset.sql_lab import cancel_query
+
+
[email protected](
+ "state", ["PENDING_STATE", "INITIALIZED_STATE", "RUNNING_STATE"]
+)
+def test_cancel_unfinished_operation(state: str) -> None:
+ """An early stop must reach the live operation in every unfinished
state."""
+ query = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0)
+ cursor = Mock()
+ cursor.status.side_effect = [state, "FINISHED_STATE"]
+ with patch("superset.db_engine_specs.impala.db") as db:
+ db.session.query.return_value.filter_by.return_value.one.return_value
= query
+ ImpalaEngineSpec.handle_cursor(cursor, query)
+ cursor.cancel_operation.assert_called_once_with()
+ cursor.close_operation.assert_called_once_with()
+ cursor.close.assert_called_once_with()
+ cursor.get_log.assert_not_called()
+
+
+def test_pending_operation_is_polled_without_progress() -> None:
+ """Pending work must not escape the cancel/progress polling loop."""
+ app = Flask(__name__)
+ app.config["DB_POLL_INTERVAL_SECONDS"] = {"impala": 0}
+ query = Mock(id=1, extra={}, progress=0)
+ cursor = Mock()
+ cursor.status.side_effect = [
+ "PENDING_STATE",
+ "INITIALIZED_STATE",
+ "RUNNING_STATE",
+ "FINISHED_STATE",
+ ]
+ cursor.get_log.return_value = "Query abc: 25% Complete"
+ with app.app_context(), patch("superset.db_engine_specs.impala.db") as db:
+ db.session.query.return_value.filter_by.return_value.one.return_value
= query
+ ImpalaEngineSpec.handle_cursor(cursor, query)
+ assert cursor.status.call_count == 4
+ cursor.get_log.assert_called_once_with()
+ assert query.progress == 25
+ cursor.cancel_operation.assert_not_called()
+
+
[email protected]("use_executor", [False, True])
+def test_stop_with_cancel_id_uses_http(use_executor: bool) -> None:
+ """Stop must reach Impala even when no worker is polling the live
cursor."""
+ app = Flask(__name__)
+ app.config["IMPALA_CANCEL_QUERY_ALLOW_INTERNAL_HOSTS"] = True
+ cancel_id = "0123456789abcdef:fedcba9876543210"
+ query = MagicMock(extra={QUERY_CANCEL_KEY: cancel_id})
Review Comment:
Annotated the mock and local variables in 1d358309c4 (`app: Flask`,
`cancel_id: str`, `query: MagicMock`, `post: MagicMock` declared before the
`with`, plus `query`/`cursor` in the other 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]
