fix: Avoid 500 if end users write bad SQL (#26638)

This commit is contained in:
Igor Khrol 2024-01-17 19:48:50 +02:00 committed by GitHub
parent eff4422d04
commit 80a6e25a98
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 44 additions and 3 deletions

View File

@ -38,4 +38,4 @@ class SupersetDBAPIOperationalError(SupersetDBAPIError):
class SupersetDBAPIProgrammingError(SupersetDBAPIError):
pass
status = 400

View File

@ -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]:
"""

View File

@ -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