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,