From 3daa038f5f6cc02b77bf7a03396fb31261d28dbd Mon Sep 17 00:00:00 2001
From: Igor Khrol <igor.khrol@automattic.com>
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 == []