From 80a6e25a98fe05f31a3c265d461c0825fa7d0aef Mon Sep 17 00:00:00 2001
From: Igor Khrol <igor.khrol@automattic.com>
Date: Wed, 17 Jan 2024 19:48:50 +0200
Subject: [PATCH] fix: Avoid 500 if end users write bad SQL (#26638)

---
 superset/db_engine_specs/exceptions.py        |  2 +-
 superset/db_engine_specs/trino.py             | 26 +++++++++++++++++--
 .../unit_tests/db_engine_specs/test_trino.py  | 19 ++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/superset/db_engine_specs/exceptions.py b/superset/db_engine_specs/exceptions.py
index 56e354d62..93a6b9927 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 6bdeae9d7..6d95f9589 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -32,7 +32,12 @@ from superset import db
 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
@@ -328,11 +333,28 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
     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 df1457dca..308902c90 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 @@ from unittest.mock import Mock, patch
 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 (
@@ -529,3 +537,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