Skip to content

Commit

Permalink
fix: Trino - handle table not found in SQLLab (#26355)
Browse files Browse the repository at this point in the history
Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>
  • Loading branch information
Khrol and john-bodley authored Jan 11, 2024
1 parent b2a21f7 commit 3daa038
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 0 deletions.
25 changes: 25 additions & 0 deletions superset/db_engine_specs/trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from flask import current_app
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import URL
from sqlalchemy.exc import NoSuchTableError
from sqlalchemy.orm import Session

from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY, USER_AGENT
Expand Down Expand Up @@ -395,3 +396,27 @@ def get_columns(
return base_cols

return [col for base_col in base_cols for col in cls._expand_columns(base_col)]

@classmethod
def get_indexes(
cls,
database: Database,
inspector: Inspector,
table_name: str,
schema: str | None,
) -> list[dict[str, Any]]:
"""
Get the indexes associated with the specified schema/table.
Trino dialect raises NoSuchTableError in get_indexes if table is empty.
:param database: The database to inspect
:param inspector: The SQLAlchemy inspector
:param table_name: The table to inspect
:param schema: The schema to inspect
:returns: The indexes
"""
try:
return super().get_indexes(database, inspector, table_name, schema)
except NoSuchTableError:
return []
16 changes: 16 additions & 0 deletions tests/unit_tests/db_engine_specs/test_trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -517,3 +517,19 @@ def test_get_columns_expand_rows(mocker: MockerFixture):
]

_assert_columns_equal(actual, expected)


def test_get_indexes_no_table():
from sqlalchemy.exc import NoSuchTableError

from superset.db_engine_specs.trino import TrinoEngineSpec

db_mock = Mock()
inspector_mock = Mock()
inspector_mock.get_indexes = Mock(
side_effect=NoSuchTableError("The specified table does not exist.")
)
result = TrinoEngineSpec.get_indexes(
db_mock, inspector_mock, "test_table", "test_schema"
)
assert result == []

0 comments on commit 3daa038

Please sign in to comment.