diff --git a/superset/db_engine_specs/exceptions.py b/superset/db_engine_specs/exceptions.py index 56e354d62a574..93a6b99271d03 100644 --- a/superset/db_engine_specs/exceptions.py +++ b/superset/db_engine_specs/exceptions.py @@ -38,4 +38,4 @@ class SupersetDBAPIOperationalError(SupersetDBAPIError): class SupersetDBAPIProgrammingError(SupersetDBAPIError): - pass + status = 400 diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 1dc711880d729..ccbb73a81c67d 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -32,7 +32,12 @@ from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import BaseEngineSpec -from superset.db_engine_specs.exceptions import SupersetDBAPIConnectionError +from superset.db_engine_specs.exceptions import ( + SupersetDBAPIConnectionError, + SupersetDBAPIDatabaseError, + SupersetDBAPIOperationalError, + SupersetDBAPIProgrammingError, +) from superset.db_engine_specs.presto import PrestoBaseEngineSpec from superset.models.sql_lab import Query from superset.superset_typing import ResultSetColumnType @@ -330,11 +335,28 @@ def update_params_from_encrypted_extra( def get_dbapi_exception_mapping(cls) -> dict[type[Exception], type[Exception]]: # pylint: disable=import-outside-toplevel from requests import exceptions as requests_exceptions + from trino import exceptions as trino_exceptions - return { + static_mapping: dict[type[Exception], type[Exception]] = { requests_exceptions.ConnectionError: SupersetDBAPIConnectionError, } + class _CustomMapping(dict[type[Exception], type[Exception]]): + def get( # type: ignore[override] + self, item: type[Exception], default: type[Exception] | None = None + ) -> type[Exception] | None: + if static := static_mapping.get(item): + return static + if issubclass(item, trino_exceptions.InternalError): + return SupersetDBAPIDatabaseError + if issubclass(item, trino_exceptions.OperationalError): + return SupersetDBAPIOperationalError + if issubclass(item, trino_exceptions.ProgrammingError): + return SupersetDBAPIProgrammingError + return default + + return _CustomMapping() + @classmethod def _expand_columns(cls, col: ResultSetColumnType) -> list[ResultSetColumnType]: """ diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py index fd553b241c8f1..751ffbcfbc5cd 100644 --- a/tests/unit_tests/db_engine_specs/test_trino.py +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -24,11 +24,19 @@ import pandas as pd import pytest from pytest_mock import MockerFixture +from requests.exceptions import ConnectionError as RequestsConnectionError from sqlalchemy import types +from trino.exceptions import TrinoExternalError, TrinoInternalError, TrinoUserError from trino.sqlalchemy import datatype import superset.config from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT +from superset.db_engine_specs.exceptions import ( + SupersetDBAPIConnectionError, + SupersetDBAPIDatabaseError, + SupersetDBAPIOperationalError, + SupersetDBAPIProgrammingError, +) from superset.superset_typing import ResultSetColumnType, SQLAColumnType from superset.utils.core import GenericDataType from tests.unit_tests.db_engine_specs.utils import ( @@ -533,3 +541,14 @@ def test_get_indexes_no_table(): db_mock, inspector_mock, "test_table", "test_schema" ) assert result == [] + + +def test_get_dbapi_exception_mapping(): + from superset.db_engine_specs.trino import TrinoEngineSpec + + mapping = TrinoEngineSpec.get_dbapi_exception_mapping() + assert mapping.get(TrinoUserError) == SupersetDBAPIProgrammingError + assert mapping.get(TrinoInternalError) == SupersetDBAPIDatabaseError + assert mapping.get(TrinoExternalError) == SupersetDBAPIOperationalError + assert mapping.get(RequestsConnectionError) == SupersetDBAPIConnectionError + assert mapping.get(Exception) is None