From 09c9410bdcf0000210e863e8f91cfef892a2102a Mon Sep 17 00:00:00 2001
From: Elizabeth Thompson <eschutho@gmail.com>
Date: Sun, 2 Aug 2026 18:18:31 +0000
Subject: [PATCH] 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              | 10 ++++++++--
 tests/unit_tests/sql_lab_test.py | 25 +++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index ae135705018..b98eb3a6e03 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -162,8 +162,14 @@ def get_query(query_id: int) -> Query:
         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()
+        # 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
 
 
diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py
index ec2745639b9..313c556a945 100644
--- a/tests/unit_tests/sql_lab_test.py
+++ b/tests/unit_tests/sql_lab_test.py
@@ -37,6 +37,7 @@ from superset.sql_lab import (
     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
@@ -100,6 +101,30 @@ def test_get_query_rolls_back_session_before_retrying(
     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,