From 7d5e8ad785c1fd1ca04919add6fdf567a3fb9a6e Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Sat, 1 Aug 2026 15:12:43 +0000 Subject: [PATCH 1/2] 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 --- superset/sql_lab.py | 3 +++ tests/unit_tests/sql_lab_test.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index bf17daac6800..ae1357050187 100644 --- a/superset/sql_lab.py +++ b/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 diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index b9bec3cad4cc..ec2745639b9a 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -35,6 +35,7 @@ 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, From 09c9410bdcf0000210e863e8f91cfef892a2102a Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Sun, 2 Aug 2026 18:18:31 +0000 Subject: [PATCH 2/2] 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 --- 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 ae1357050187..b98eb3a6e031 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 ec2745639b9a..313c556a945c 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -37,6 +37,7 @@ 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,