Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Jan 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion superset/db_engine_specs/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,4 @@ class SupersetDBAPIOperationalError(SupersetDBAPIError):


class SupersetDBAPIProgrammingError(SupersetDBAPIError):
pass
status = 400
26 changes: 24 additions & 2 deletions superset/db_engine_specs/trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]:
"""
Expand Down
19 changes: 19 additions & 0 deletions tests/unit_tests/db_engine_specs/test_trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
Loading