fix(api): handle DBAPI ProgrammingError on database connect in schemas and chart data

superset/charts/data/api.py

@@ -27,6 +27,7 @@
 from flask_babel import gettext as _
 from flask_caching.backends import NullCache
 from marshmallow import ValidationError
+from sqlalchemy.exc import DBAPIError
 from werkzeug.utils import secure_filename
 
 from superset import is_feature_enabled, security_manager
@@ -54,6 +55,7 @@
 from superset.connectors.sqla.models import BaseDatasource
 from superset.constants import CACHE_DISABLED_TIMEOUT, SemanticCacheStatus
 from superset.daos.exceptions import DatasourceNotFound
+from superset.databases.error_provenance import is_database_engine_error
 from superset.errors import SupersetErrorType
 from superset.exceptions import QueryObjectValidationError, SupersetSecurityException
 from superset.extensions import cache_manager, event_logger
@@ -699,6 +701,21 @@ def _get_data_response(
             QueryObjectValidationError,
         ) as exc:
             return self.response_400(message=sanitize_error_message(exc.message))
+        except DBAPIError as ex:
+            if not is_database_engine_error(ex):
+                # e.g. a Superset metadata-DB failure: keep it a 5xx
+                raise
+            # Query execution failures are wrapped as ChartDataQueryFailedError, so a
+            # raw DBAPIError here comes from acquiring a connection outside query
+            # execution (e.g. dialect initialization while computing the cache key).
+            logger.warning(
+                "Database error while acquiring chart data payload", exc_info=True
+            )
+            return self.response_422(
+                message=sanitize_error_message(
+                    _("There was an error connecting to the database")
+                )
+            )
 
             # Log is_cached if extra payload callback is provided
         materialized_result = result.materialize()

superset/databases/api.py

@@ -36,7 +36,7 @@
 from flask_appbuilder.api.schemas import get_list_schema
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from marshmallow import ValidationError
-from sqlalchemy.exc import NoSuchTableError, OperationalError, SQLAlchemyError
+from sqlalchemy.exc import DBAPIError, NoSuchTableError, SQLAlchemyError
 
 from superset import event_logger
 from superset.commands.database.create import CreateDatabaseCommand
@@ -89,6 +89,7 @@
 from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
 from superset.daos.database import DatabaseDAO
 from superset.databases.decorators import check_table_access
+from superset.databases.error_provenance import is_database_engine_error
 from superset.databases.filters import (
     DatabaseFilter,
     DatabaseOAuth2TokensUserFilter,
@@ -782,9 +783,14 @@ def catalogs(self, pk: int, **kwargs: Any) -> FlaskResponse:
                 catalogs,
             )
             return self.response(200, result=list(catalogs))
-        except OperationalError:
+        except DBAPIError as ex:
+            if not is_database_engine_error(ex):
+                raise
+            logger.warning(
+                "Error connecting to database %s to list catalogs", pk, exc_info=True
+            )
             return self.response(
-                500,
+                422,
                 message="There was an error connecting to the database",
             )
         except OAuth2RedirectError:
@@ -867,9 +873,14 @@ def schemas(self, pk: int, **kwargs: Any) -> FlaskResponse:
                         ],
                     )
             return self.response(200, result=list(schemas))
-        except OperationalError:
+        except DBAPIError as ex:
+            if not is_database_engine_error(ex):
+                raise
+            logger.warning(
+                "Error connecting to database %s to list schemas", pk, exc_info=True
+            )
             return self.response(
-                500, message="There was an error connecting to the database"
+                422, message="There was an error connecting to the database"
             )
         except OAuth2RedirectError:
             raise

superset/translations/messages.pot

@@ -14995,6 +14995,9 @@ msgstr ""
 msgid "There was a problem refreshing your dashboard. We'll try again in %s, as scheduled."
 msgstr ""
 
+msgid "There was an error connecting to the database"
+msgstr ""
+
 msgid "There was an error creating the group. Please, try again."
 msgstr ""
 

tests/unit_tests/charts/test_chart_data_api.py

@@ -1665,3 +1665,46 @@ def test_get_data_response_streaming_up_front_maps_setup_errors(
     assert response.json["message"] == str(exc)
     mock_stream.assert_called_once()
     command.execute.assert_not_called()
+
+
+def test_get_data_response_handles_dbapi_error_on_connect(app: SupersetApp) -> None:
+    """
+    A raw DBAPIError raised outside query execution (e.g. while computing the cache
+    key) returns a generic 422 instead of escaping as a 500.
+    """
+    from sqlalchemy.exc import ProgrammingError
+
+    from superset.databases.error_provenance import _DATABASE_ENGINE_ERROR_MARKER
+
+    error = ProgrammingError(
+        "show standard_conforming_strings", {}, Exception("secret driver detail")
+    )
+    setattr(error, _DATABASE_ENGINE_ERROR_MARKER, True)  # noqa: B010
+    command = MagicMock()
+    command.execute.side_effect = error
+    api = ChartDataRestApi()
+
+    with app.test_request_context("/api/v1/chart/data"):
+        response = api._get_data_response(command)
+
+    assert response.status_code == 422
+    body = response.get_data(as_text=True)
+    assert "There was an error connecting to the database" in body
+    assert "secret driver detail" not in body
+    assert "standard_conforming_strings" not in body
+
+
+def test_get_data_response_reraises_unmarked_dbapi_error(app: SupersetApp) -> None:
+    """
+    A DBAPIError that did not come from a customer Database engine (e.g. the
+    Superset metadata DB) is not reported as a connection error to the user.
+    """
+    from sqlalchemy.exc import ProgrammingError
+
+    command = MagicMock()
+    command.execute.side_effect = ProgrammingError("select 1", {}, Exception("meta"))
+    api = ChartDataRestApi()
+
+    with app.test_request_context("/api/v1/chart/data"):
+        with pytest.raises(ProgrammingError):
+            api._get_data_response(command)

tests/unit_tests/databases/api_test.py

@@ -3044,3 +3044,72 @@ def test_related_objects_limits_dataset_details(
     assert payload["count"] == MAX_RELATED_DATASETS + 2
     assert len(payload["result"]) == MAX_RELATED_DATASETS
     assert can_access.call_count == MAX_RELATED_DATASETS
+
+
+@pytest.mark.parametrize(
+    "endpoint, method",
+    [("schemas", "get_all_schema_names"), ("catalogs", "get_all_catalog_names")],
+)
+def test_schemas_and_catalogs_dbapi_programming_error(
+    mocker: MockerFixture,
+    client: Any,
+    full_api_access: None,
+    endpoint: str,
+    method: str,
+) -> None:
+    """
+    A `ProgrammingError` raised while connecting (e.g. dialect initialization
+    against a Postgres-compatible engine) is reported like other connection errors.
+    """
+    from sqlalchemy.exc import ProgrammingError
+
+    from superset.databases.api import DatabaseRestApi
+    from superset.databases.error_provenance import _DATABASE_ENGINE_ERROR_MARKER
+
+    error = ProgrammingError(
+        "show standard_conforming_strings", {}, Exception("unrecognized parameter")
+    )
+    setattr(error, _DATABASE_ENGINE_ERROR_MARKER, True)  # noqa: B010
+    database = mocker.MagicMock()
+    getattr(database, method).side_effect = error
+    DatabaseDAO = mocker.patch("superset.databases.api.DatabaseDAO")  # noqa: N806
+    DatabaseDAO.find_by_id.return_value = database
+    mocker.patch.object(DatabaseRestApi, "datamodel").get.return_value = database
+
+    response = client.get(f"/api/v1/database/1/{endpoint}/")
+
+    assert response.status_code == 422
+    assert response.json == {"message": "There was an error connecting to the database"}
+
+
+@pytest.mark.parametrize(
+    "endpoint, method",
+    [("schemas", "get_all_schema_names"), ("catalogs", "get_all_catalog_names")],
+)
+def test_schemas_and_catalogs_unmarked_dbapi_error_not_masked(
+    mocker: MockerFixture,
+    client: Any,
+    full_api_access: None,
+    endpoint: str,
+    method: str,
+) -> None:
+    """
+    A DBAPIError not marked as coming from a Database engine (e.g. a Superset
+    metadata-DB failure) is not reported as a customer connection error.
+    """
+    from sqlalchemy.exc import ProgrammingError
+
+    from superset.databases.api import DatabaseRestApi
+
+    database = mocker.MagicMock()
+    getattr(database, method).side_effect = ProgrammingError(
+        "select 1", {}, Exception("metadata failure")
+    )
+    DatabaseDAO = mocker.patch("superset.databases.api.DatabaseDAO")  # noqa: N806
+    DatabaseDAO.find_by_id.return_value = database
+    mocker.patch.object(DatabaseRestApi, "datamodel").get.return_value = database
+
+    response = client.get(f"/api/v1/database/1/{endpoint}/")
+
+    assert response.status_code == 500
+    assert "error connecting to the database" not in response.get_data(as_text=True)