fix(sqllab): don't swallow OAuth2RedirectError in query cost estimation

The broad except-Exception clause added to wrap raw DBAPI errors also
caught OAuth2RedirectError (raised via get_raw_connection -> check_for_
oauth2), re-wrapping it as SupersetGenericDBErrorException (400) and
discarding the redirect metadata the frontend needs to drive interactive
re-auth. Add an except OAuth2RedirectError: raise clause before the
catch-all, mirroring the existing guard in sql_lab.py's
execute_sql_statements().

Addresses CHANGES_REQUESTED review feedback from rebenitez1802.

superset/commands/sql_lab/estimate.py

@@ -28,10 +28,12 @@ from superset.commands.base import BaseCommand
 from superset.daos.database import DatabaseDAO
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import (
+    OAuth2RedirectError,
     SupersetDisallowedSQLFunctionException,
     SupersetDisallowedSQLTableException,
     SupersetDMLNotAllowedException,
     SupersetErrorException,
+    SupersetGenericDBErrorException,
     SupersetTimeoutException,
 )
 from superset.jinja_context import get_template_processor
@@ -209,6 +211,14 @@ class QueryEstimationCommand(BaseCommand):
                 ),
                 status=500,
             ) from ex
+        except OAuth2RedirectError:
+            # user needs to authenticate with OAuth2 in order to run query
+            raise
+        except Exception as ex:
+            logger.exception("Query cost estimation failed unexpectedly")
+            raise SupersetGenericDBErrorException(
+                utils.error_msg_from_exception(ex)
+            ) from ex
 
         spec = self._database.db_engine_spec
         query_cost_formatters: dict[str, Any] = app.config[

tests/unit_tests/commands/sql_lab/test_estimate.py

@@ -26,7 +26,12 @@ from superset.commands.sql_lab.estimate import (
     QueryEstimationCommand,
 )
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import SupersetErrorException, SupersetSecurityException
+from superset.exceptions import (
+    OAuth2RedirectError,
+    SupersetErrorException,
+    SupersetGenericDBErrorException,
+    SupersetSecurityException,
+)
 
 
 def _make_params(**kwargs: object) -> EstimateQueryCostType:
@@ -445,3 +450,88 @@ def test_run_wraps_raw_jinja_undefined_error(
 
     assert exc_info.value.status == 400
     assert exc_info.value.error.error_type == SupersetErrorType.GENERIC_COMMAND_ERROR
+
+
+# ---------------------------------------------------------------------------
+# Raw DBAPI errors from estimate_query_cost must not leak from run()
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.security_manager", new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_wraps_raw_dbapi_error_from_cost_estimation(
+    mock_dao: MagicMock,
+    mock_security_manager: MagicMock,
+    mock_app: MagicMock,
+) -> None:
+    """A raw DBAPI exception (e.g. ``psycopg2.errors.SyntaxError``) raised by
+    ``estimate_query_cost`` when EXPLAIN parses invalid SQL must not leak past
+    ``run()`` — it should surface as ``SupersetGenericDBErrorException`` with
+    ``status == 400`` and ``GENERIC_DB_ENGINE_ERROR``, mirroring the sibling
+    pattern in ``SynchronousSqlJsonExecutor.execute()``."""
+    import psycopg2.errors
+
+    mock_database = MagicMock()
+    mock_dao.find_by_id.return_value = mock_database
+    mock_security_manager.raise_for_access.return_value = None
+    mock_app.config = {
+        "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+        "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+        "DISALLOWED_SQL_FUNCTIONS": {},
+        "DISALLOWED_SQL_TABLES": {},
+    }
+    mock_database.db_engine_spec.estimate_query_cost.side_effect = (
+        psycopg2.errors.UndefinedTable('relation "nonexistent_table" does not exist')
+    )
+
+    sql = "SELECT 1 FROM nonexistent_table"
+    command = QueryEstimationCommand(_make_params(sql=sql))
+    with pytest.raises(SupersetGenericDBErrorException) as exc_info:
+        command.run()
+
+    assert exc_info.value.status == 400
+    assert exc_info.value.error.error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+
+
+# ---------------------------------------------------------------------------
+# OAuth2RedirectError must pass through run() untouched, not be swallowed
+# by the broad DBAPI-error catch-all above
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.commands.sql_lab.estimate.app")
+@patch("superset.commands.sql_lab.estimate.security_manager", new_callable=MagicMock)
+@patch("superset.commands.sql_lab.estimate.DatabaseDAO")
+def test_run_reraises_oauth2_redirect_error_from_cost_estimation(
+    mock_dao: MagicMock,
+    mock_security_manager: MagicMock,
+    mock_app: MagicMock,
+) -> None:
+    """``OAuth2RedirectError`` raised by ``estimate_query_cost`` (via
+    ``get_raw_connection`` -> ``check_for_oauth2``) must propagate unchanged
+    so the frontend can drive the interactive re-auth flow — it must not be
+    re-wrapped into ``SupersetGenericDBErrorException`` by the broad
+    except-Exception clause, mirroring the sibling guard in
+    ``sql_lab.py``'s ``execute_sql_statements()``."""
+    mock_database = MagicMock()
+    mock_dao.find_by_id.return_value = mock_database
+    mock_security_manager.raise_for_access.return_value = None
+    mock_app.config = {
+        "SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT": 10,
+        "QUERY_COST_FORMATTERS_BY_ENGINE": {},
+        "DISALLOWED_SQL_FUNCTIONS": {},
+        "DISALLOWED_SQL_TABLES": {},
+    }
+    mock_database.db_engine_spec.estimate_query_cost.side_effect = OAuth2RedirectError(
+        url="https://example.org/oauth2/authorize",
+        tab_id="tab-1",
+        redirect_uri="https://example.org/oauth2/callback",
+    )
+
+    sql = "SELECT 1"
+    command = QueryEstimationCommand(_make_params(sql=sql))
+    with pytest.raises(OAuth2RedirectError) as exc_info:
+        command.run()
+
+    assert exc_info.value.status == 403