fix(sqllab): roll back session before retrying get_query after a broken transaction

`get_query` catches any exception from the ORM lookup and relies on the
`backoff` decorator to retry up to 5 times. When the underlying failure is
(or causes) a SQLAlchemy `PendingRollbackError` - e.g. a `PendingRollbackError`
chained under an `OperationalError`/`QueryCanceled` from a dropped connection
or statement timeout - the session is left in a broken state that SQLAlchemy
refuses to use again until `.rollback()` is called explicitly. Since the
session was never rolled back, every one of the 5 retries reused the same
poisoned session and failed identically, so the retry loop never had a
chance to recover from what may be a transient connection blip.

Roll back the session in the except block before raising `SqlLabException`
so each `backoff` retry starts from a clean session. The exception raised
and logged is unchanged.

Fixes SUPERSET-PYTHON-WDZ

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

superset/sql_lab.py

@@ -161,6 +161,9 @@ 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
+        db.session.rollback()
         raise SqlLabException("Failed at getting query") from ex
 
 

tests/unit_tests/sql_lab_test.py

@@ -35,6 +35,7 @@ from superset.sql.parse import SQLStatement, Table
 from superset.sql_lab import (
     execute_query,
     execute_sql_statements,
+    get_query,
     get_sql_results,
 )
 from superset.utils.rls import apply_rls, get_predicates_for_table
@@ -71,6 +72,34 @@ 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()
+
+
 @with_config(
     {
         "SQLLAB_PAYLOAD_MAX_MB": 50,