From 3daa038f5f6cc02b77bf7a03396fb31261d28dbd Mon Sep 17 00:00:00 2001 From: Igor Khrol Date: Thu, 11 Jan 2024 02:37:18 +0200 Subject: [PATCH] fix: Trino - handle table not found in SQLLab (#26355) Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com> --- superset/db_engine_specs/trino.py | 25 +++++++++++++++++++ .../unit_tests/db_engine_specs/test_trino.py | 16 ++++++++++++ 2 files changed, 41 insertions(+) diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 6e56dbfa24d6b..1dc711880d729 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -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 @@ -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 [] diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py index 15e55fc5af62f..fd553b241c8f1 100644 --- a/tests/unit_tests/db_engine_specs/test_trino.py +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -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 == []