I3eka commented on code in PR #45561:
URL: https://github.com/apache/superset/pull/45561#discussion_r4384233632
##########
tests/unit_tests/sql/execution/test_executor.py:
##########
@@ -763,6 +764,209 @@ def test_execute_custom_timeout(
)
[email protected]("failure", ["database", "timeout", "oauth"])
+def test_execute_failure_finishes_query_record(
+ mocker: MockerFixture,
+ database: Database,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_db_session: MagicMock,
+ default_sql_config: None,
+ failure: str,
+) -> None:
+ """An execution failure must not leave its audit record running."""
+ from superset.errors import ErrorLevel, SupersetErrorType
+ from superset.exceptions import OAuth2Error, SupersetTimeoutException
+ from superset.sql.execution.executor import SQLExecutor
+
+ errors = {
+ "database": RuntimeError("Query failed"),
+ "timeout": SupersetTimeoutException(
+ error_type=SupersetErrorType.GENERIC_BACKEND_ERROR,
+ message="Query timed out",
+ level=ErrorLevel.ERROR,
+ ),
+ "oauth": OAuth2Error("Authentication required"),
+ }
+ mock_query.status = "running"
+ mocker.patch.object(SQLExecutor, "_create_query_record",
return_value=mock_query)
+ mocker.patch.object(SQLExecutor, "_try_get_cached_result",
return_value=None)
+ mocker.patch.object(SQLExecutor, "_execute_statements",
side_effect=errors[failure])
+ mocker.patch(
+ "superset.sql.execution.executor.now_as_float",
+ return_value=1234000.0,
+ create=True,
+ )
Review Comment:
Rechecked on c8c079654469: `create=True` is already removed from this patch
target. The regression now requires the imported `now_as_float` attribute to
exist, and passes in the refreshed executor suite.
##########
tests/unit_tests/sql/execution/test_executor.py:
##########
@@ -763,6 +764,209 @@ def test_execute_custom_timeout(
)
[email protected]("failure", ["database", "timeout", "oauth"])
+def test_execute_failure_finishes_query_record(
+ mocker: MockerFixture,
+ database: Database,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_db_session: MagicMock,
+ default_sql_config: None,
+ failure: str,
+) -> None:
+ """An execution failure must not leave its audit record running."""
+ from superset.errors import ErrorLevel, SupersetErrorType
+ from superset.exceptions import OAuth2Error, SupersetTimeoutException
+ from superset.sql.execution.executor import SQLExecutor
+
+ errors = {
+ "database": RuntimeError("Query failed"),
+ "timeout": SupersetTimeoutException(
+ error_type=SupersetErrorType.GENERIC_BACKEND_ERROR,
+ message="Query timed out",
+ level=ErrorLevel.ERROR,
+ ),
+ "oauth": OAuth2Error("Authentication required"),
+ }
+ mock_query.status = "running"
+ mocker.patch.object(SQLExecutor, "_create_query_record",
return_value=mock_query)
+ mocker.patch.object(SQLExecutor, "_try_get_cached_result",
return_value=None)
+ mocker.patch.object(SQLExecutor, "_execute_statements",
side_effect=errors[failure])
+ mocker.patch(
+ "superset.sql.execution.executor.now_as_float",
+ return_value=1234000.0,
+ create=True,
+ )
+ expected_status = (
+ QueryStatus.TIMED_OUT if failure == "timeout" else QueryStatus.FAILED
+ )
+
+ if failure == "oauth":
+ with pytest.raises(OAuth2Error):
+ database.execute("SELECT 1")
+ else:
+ result = database.execute("SELECT 1")
+ assert result.status == expected_status
+ assert result.query_id == mock_query.id
+ assert result.error_message == mock_query.error_message
+
+ assert mock_query.status == expected_status.value
+ assert mock_query.error_message
+ assert mock_query.end_time == 1234000.0
+ mock_db_session.rollback.assert_called_once()
+ mock_db_session.commit.assert_called_once()
+
+
[email protected]("failure", ["database", "timeout", "oauth",
"oauth_redirect"])
[email protected]("metadata_failure", ["rollback", "commit"])
+def test_execute_preserves_error_when_failure_recording_fails(
+ mocker: MockerFixture,
+ database: Database,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_db_session: MagicMock,
+ default_sql_config: None,
+ failure: str,
+ metadata_failure: str,
+ caplog: pytest.LogCaptureFixture,
+) -> None:
+ """A metadata outage must not replace the execution error or OAuth
redirect."""
+ from superset.errors import ErrorLevel, SupersetErrorType
+ from superset.exceptions import (
+ OAuth2Error,
+ OAuth2RedirectError,
+ SupersetTimeoutException,
+ )
+ from superset.sql.execution.executor import SQLExecutor
+
+ error = {
+ "database": RuntimeError("Query failed"),
+ "timeout": SupersetTimeoutException(
+ error_type=SupersetErrorType.GENERIC_BACKEND_ERROR,
+ message="Query timed out",
+ level=ErrorLevel.ERROR,
+ ),
+ "oauth": OAuth2Error("Authentication required"),
+ "oauth_redirect": OAuth2RedirectError(
+ "https://example.com/auth";, "tab-id",
"https://example.com/callback";
+ ),
+ }[failure]
+ mocker.patch.object(SQLExecutor, "_create_query_record",
return_value=mock_query)
+ mocker.patch.object(SQLExecutor, "_try_get_cached_result",
return_value=None)
+ mocker.patch.object(SQLExecutor, "_execute_statements", side_effect=error)
+ getattr(mock_db_session, metadata_failure).side_effect = RuntimeError(
+ "Metadata unavailable"
+ )
+
+ if isinstance(error, (OAuth2Error, OAuth2RedirectError)):
+ with pytest.raises(type(error)) as caught:
+ database.execute("SELECT 1")
+ assert caught.value is error
+ else:
+ result = database.execute("SELECT 1")
+ assert result.status == (
+ QueryStatus.TIMED_OUT if failure == "timeout" else
QueryStatus.FAILED
+ )
+ assert result.error_message == (
+ "Query exceeded the timeout limit"
+ if failure == "timeout"
+ else "sqlite error: Query failed"
+ )
+ if metadata_failure == "commit":
+ assert result.query_id == mock_query.id
+
+ assert mock_db_session.rollback.call_count == 2
+ assert "Failed to persist SQL query error" in caplog.text
Review Comment:
Rechecked on c8c079654469: the test already calls
`caplog.set_level("WARNING", logger="superset.sql.execution.executor")` before
execution. All eight commit/rollback and execution-error combinations pass and
assert the warning content.
##########
superset/sql/execution/executor.py:
##########
@@ -577,6 +579,7 @@ def execute(
opts: QueryOptionsType = options or QueryOptionsType()
start_time = time.time()
+ query: Query | None = None
Review Comment:
This is a false positive on the refreshed c8c079654469 head. Postponed
annotations are enabled at the top of the module. A fresh-process import
succeeds without `Query` in module globals, and the return and parameter
annotations remain `"Query"` and `"Query | None"` strings.
##########
superset/sql/execution/executor.py:
##########
@@ -1166,7 +1204,7 @@ def _create_query_record(
catalog: str | None,
schema: str | None,
status: QueryStatus,
- ) -> Any:
+ ) -> Query:
Review Comment:
No runtime import is needed here: c8c079654469 already enables postponed
annotations. I verified a fresh import and the string annotations for both
affected methods, then ran the executor tests. Adding a runtime `Query` import
would not fix an actual failure.
##########
superset/sql/execution/executor.py:
##########
@@ -1008,14 +1020,40 @@ def _create_error_result(
:param start_time: Start time for calculating execution duration
:param partial_results: Optional list of StatementResult from
successful
statements before the failure
+ :param query: Persisted audit record, if execution reached that stage
:returns: QueryResult with error status
"""
- from superset_core.queries.types import QueryResult as QueryResultType
+ from superset_core.queries.types import (
+ QueryResult as QueryResultType,
+ QueryStatus,
+ )
+
+ from superset.sql.execution.utils import commit_query_unless_stopped
+
+ query_id = None
+ if query is not None:
+ try:
+ # A failed metadata flush must be rolled back before recording
failure.
+ if not db.session.is_active:
+ db.session.rollback() # pylint:
disable=consider-using-transaction
+ query_id = query.id
+ query.status = status.value
+ query.error_message = error_message
+ query.end_time = now_as_float()
Review Comment:
Rechecked on c8c079654469. The persisted identity is captured with
`inspect(query).identity[0]` before rollback or expired-attribute loading. The
two real-session regressions preserve query ID 123 when rollback fails or
reloading an expired attribute fails; both pass.
##########
superset/sql/execution/executor.py:
##########
@@ -1008,14 +1020,40 @@ def _create_error_result(
:param start_time: Start time for calculating execution duration
:param partial_results: Optional list of StatementResult from
successful
statements before the failure
+ :param query: Persisted audit record, if execution reached that stage
:returns: QueryResult with error status
"""
- from superset_core.queries.types import QueryResult as QueryResultType
+ from superset_core.queries.types import (
+ QueryResult as QueryResultType,
+ QueryStatus,
+ )
+
+ from superset.sql.execution.utils import commit_query_unless_stopped
+
+ query_id = None
+ if query is not None:
+ try:
+ # A failed metadata flush must be rolled back before recording
failure.
+ if not db.session.is_active:
+ db.session.rollback() # pylint:
disable=consider-using-transaction
+ query_id = query.id
+ query.status = status.value
+ query.error_message = error_message
+ query.end_time = now_as_float()
+ if not commit_query_unless_stopped(query):
+ status = QueryStatus.STOPPED
+ error_message = query.error_message
Review Comment:
Rechecked the shared helper on c8c079654469. Its False path already
refreshes the Query with `with_for_update=True` inside `no_autoflush`,
replacing the staged failure fields with the stored Stop. The caller commits
after reading that result. The regression verifies that a later caller commit
cannot flush stale failure metadata, so another refresh is unnecessary.
--
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]
