Copilot commented on code in PR #45561:
URL: https://github.com/apache/superset/pull/45561#discussion_r4335691249


##########
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:
   `Query` is referenced in a runtime-executed variable annotation, but in this 
diff `Query` is only imported under `TYPE_CHECKING`. If this module does not 
use postponed evaluation of annotations (`from __future__ import annotations`), 
this line will raise `NameError` at runtime. Fix by either importing `Query` at 
runtime (not only under `TYPE_CHECKING`) or by changing the annotation to a 
quoted forward reference (e.g., `query: "Query | None"` or `query: "Query" | 
None` depending on the project’s typing conventions).



##########
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:
   This assertion depends on log capture configuration. To make the test robust 
across different pytest logging setups, explicitly set the capture level/logger 
before executing (e.g., `caplog.set_level(logging.WARNING, 
logger="superset.sql.execution.executor")`) so the warning is always recorded.



##########
superset/sql/execution/executor.py:
##########
@@ -1008,14 +1020,39 @@ 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.
+                db.session.rollback()  # pylint: 
disable=consider-using-transaction
+                query_id = query.id

Review Comment:
   `_create_error_result` performs an unconditional `db.session.rollback()` 
when a `query` is present. In a scoped-session (request) context, this can 
unintentionally roll back unrelated pending metadata changes in the same 
session, not just the query record. Consider limiting rollback to cases where 
the session is in a failed/invalid state (e.g., after a flush/commit error), or 
isolating failure finalization into a separate/nested transaction/session so 
the error path doesn’t implicitly clear broader unit-of-work state.



##########
superset/sql/execution/executor.py:
##########
@@ -1166,7 +1203,7 @@ def _create_query_record(
         catalog: str | None,
         schema: str | None,
         status: QueryStatus,
-    ) -> Any:
+    ) -> Query:

Review Comment:
   The return annotation `-> Query` has the same runtime risk as the new local 
annotation: `Query` is imported only under `TYPE_CHECKING` in this diff. If 
postponed evaluation of annotations is not enabled for the module, this will 
raise `NameError` when defining the function. Use a quoted forward reference 
(`-> "Query"`) or import `Query` at runtime.



##########
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:
   `create=True` can mask patch-target typos by silently creating a new 
attribute if the target path/name is wrong. Since `now_as_float` is imported by 
the module in this PR, it should exist—dropping `create=True` makes the test 
fail fast if the patch path ever becomes invalid.



-- 
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]

Reply via email to