address review feedback: don't let a rollback failure escape retry contract

get_query's backoff decorator only retries on SqlLabException. If
db.session.rollback() itself raises (e.g. the connection is fully
dead), that new exception would replace the intended SqlLabException
and bypass the retry contract. Swallow rollback failures so the
original lookup error is always what gets raised.

Co-Authored-By: Claude <noreply@anthropic.com>

superset/sql_lab.py

@@ -161,6 +161,15 @@ def get_query(query_id: int) -> Query:
     try:
         return db.session.query(Query).filter_by(id=query_id).one()
     except Exception as ex:
+        # roll back so a poisoned session (e.g. PendingRollbackError after a
+        # failed flush) doesn't fail every subsequent backoff retry identically.
+        # Swallow rollback failures so a session/connection too broken to roll
+        # back doesn't replace the original exception with one the backoff
+        # decorator won't retry on.
+        try:
+            db.session.rollback()
+        except Exception:  # pylint: disable=broad-except
+            logger.warning("Failed to roll back session in get_query", exc_info=True)
         raise SqlLabException("Failed at getting query") from ex
 
 

tests/unit_tests/sql_lab_test.py

@@ -35,7 +35,9 @@ from superset.sql.parse import SQLStatement, Table
 from superset.sql_lab import (
     execute_query,
     execute_sql_statements,
+    get_query,
     get_sql_results,
+    SqlLabException,
 )
 from superset.utils.rls import apply_rls, get_predicates_for_table
 from tests.conftest import with_config
@@ -71,6 +73,58 @@ def test_execute_query(mocker: MockerFixture, app: None) -> None:
     SupersetResultSet.assert_called_with([(42,)], cursor.description, db_engine_spec)
 
 
+def test_get_query_rolls_back_session_before_retrying(
+    mocker: MockerFixture, app: SupersetApp
+) -> None:
+    """
+    A broken transaction (e.g. `PendingRollbackError` following a failed flush)
+    leaves the session unusable until `session.rollback()` is called, so without
+    it every `backoff` retry would reuse the same poisoned session and fail
+    identically. `get_query` must roll back on failure so each retry gets a
+    clean session and has a real chance to succeed.
+    """
+    # avoid actually sleeping through the `backoff` decorator's retry interval
+    mocker.patch("backoff._sync.time.sleep")
+
+    expected_query = mocker.MagicMock()
+    mock_one = mocker.patch("superset.sql_lab.db.session.query")
+    mock_one.return_value.filter_by.return_value.one.side_effect = [
+        Exception("session is broken"),
+        expected_query,
+    ]
+    mock_rollback = mocker.patch("superset.sql_lab.db.session.rollback")
+
+    result = get_query(query_id=1)
+
+    assert result is expected_query
+    assert mock_one.return_value.filter_by.return_value.one.call_count == 2
+    mock_rollback.assert_called_once()
+
+
+def test_get_query_swallows_rollback_failure(
+    mocker: MockerFixture, app: SupersetApp
+) -> None:
+    """
+    If the session/connection is too broken for `rollback()` itself to succeed,
+    that failure must not replace the original lookup error: `get_query` still
+    needs to raise `SqlLabException` so the `backoff` decorator's retry contract
+    (which only matches on `SqlLabException`) isn't bypassed.
+    """
+    mocker.patch("backoff._sync.time.sleep")
+
+    mock_one = mocker.patch("superset.sql_lab.db.session.query")
+    mock_one.return_value.filter_by.return_value.one.side_effect = Exception(
+        "session is broken"
+    )
+    mocker.patch(
+        "superset.sql_lab.db.session.rollback",
+        side_effect=Exception("connection already closed"),
+    )
+
+    with pytest.raises(SqlLabException):
+        get_query(query_id=1)
+
+
 @with_config(
     {
         "SQLLAB_PAYLOAD_MAX_MB": 50,