fix: Trino - handle table not found in SQLLab (#26355)

Co-authored-by: John Bodley <4567245+john-bodley@users.noreply.github.com>
This commit is contained in:
Igor Khrol 2024-01-11 02:37:18 +02:00 committed by GitHub
parent b2a21f7166
commit 3daa038f5f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 0 deletions

View File

@ -26,6 +26,7 @@ import simplejson as json
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 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
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 []

View File

@ -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 == []