From 12287cdd8f63f36edb81a57c7a12e3b0157ad1a1 Mon Sep 17 00:00:00 2001 From: Igor Khrol Date: Mon, 25 Dec 2023 18:01:18 +0200 Subject: [PATCH 1/5] Trino: handle table not found in SQLLab --- superset/databases/api.py | 9 ++++++--- tests/integration_tests/databases/api_tests.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/superset/databases/api.py b/superset/databases/api.py index 8de84a16af6dc..a4cd5655f4716 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -804,9 +804,12 @@ def table_extra_metadata( parsed_schema = parse_js_uri_path_item(schema_name, eval_undefined=True) table_name = cast(str, parse_js_uri_path_item(table_name)) - payload = database.db_engine_spec.extra_table_metadata( - database, table_name, parsed_schema - ) + try: + payload = database.db_engine_spec.extra_table_metadata( + database, table_name, parsed_schema + ) + except NoSuchTableError: + return self.response_404() return self.response(200, **payload) @expose("//select_star//", methods=("GET",)) diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index 0bc1f245a1f7b..23be07de4f200 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -1340,6 +1340,21 @@ def test_get_invalid_table_table_extra_metadata(self): self.assertEqual(rv.status_code, 200) self.assertEqual(data, {}) + def test_get_invalid_table_table_extra_metadata_trino(self): + """ + Database API: Test get invalid table from table extra metadata + """ + from superset.utils.database import get_or_create_db + + trino_db = get_or_create_db("trino_examples", "trino://42/sd") + uri = f"api/v1/database/{trino_db.id}/table_extra/wrong_table/null/" + self.login(username="admin") + rv = self.client.get(uri) + data = json.loads(rv.data.decode("utf-8")) + + self.assertEqual(rv.status_code, 404) + self.assertEqual(data, {"message": "Not found"}) + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_get_select_star(self): """ From 123d1f8d7ee437db88f99a8f78e4fa90618cfb56 Mon Sep 17 00:00:00 2001 From: Igor Khrol Date: Tue, 9 Jan 2024 14:34:13 +0200 Subject: [PATCH 2/5] Handle non-existing table on `extra_table_metadata` in TrinoEngineSpec --- superset/databases/api.py | 9 +++------ superset/db_engine_specs/trino.py | 5 ++++- tests/integration_tests/databases/api_tests.py | 5 +++-- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/superset/databases/api.py b/superset/databases/api.py index a4cd5655f4716..8de84a16af6dc 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -804,12 +804,9 @@ def table_extra_metadata( parsed_schema = parse_js_uri_path_item(schema_name, eval_undefined=True) table_name = cast(str, parse_js_uri_path_item(table_name)) - try: - payload = database.db_engine_spec.extra_table_metadata( - database, table_name, parsed_schema - ) - except NoSuchTableError: - return self.response_404() + payload = database.db_engine_spec.extra_table_metadata( + database, table_name, parsed_schema + ) return self.response(200, **payload) @expose("//select_star//", methods=("GET",)) diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 6e56dbfa24d6b..aa71067270cdc 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -58,7 +58,10 @@ def extra_table_metadata( table_name: str, schema_name: str | None, ) -> dict[str, Any]: - metadata = {} + metadata: dict[str, Any] = {} + + if not database.has_table_by_name(table_name): + return metadata if indexes := database.get_indexes(table_name, schema_name): col_names, latest_parts = cls.latest_partition( diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index 23be07de4f200..05713c2e8d2eb 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -1343,6 +1343,7 @@ def test_get_invalid_table_table_extra_metadata(self): def test_get_invalid_table_table_extra_metadata_trino(self): """ Database API: Test get invalid table from table extra metadata + Trino dialect """ from superset.utils.database import get_or_create_db @@ -1352,8 +1353,8 @@ def test_get_invalid_table_table_extra_metadata_trino(self): rv = self.client.get(uri) data = json.loads(rv.data.decode("utf-8")) - self.assertEqual(rv.status_code, 404) - self.assertEqual(data, {"message": "Not found"}) + self.assertEqual(rv.status_code, 200) + self.assertEqual(data, {}) @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_get_select_star(self): From 87d4ab532ef0d1e62732728fcccb8b32c5144622 Mon Sep 17 00:00:00 2001 From: Igor Khrol Date: Wed, 10 Jan 2024 16:37:32 +0200 Subject: [PATCH 3/5] Overwrite get_indexes --- superset/db_engine_specs/trino.py | 27 ++++++++++++++++--- .../unit_tests/db_engine_specs/test_trino.py | 16 +++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index aa71067270cdc..bec0bd8d604a5 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 @@ -60,9 +61,6 @@ def extra_table_metadata( ) -> dict[str, Any]: metadata: dict[str, Any] = {} - if not database.has_table_by_name(table_name): - return metadata - if indexes := database.get_indexes(table_name, schema_name): col_names, latest_parts = cls.latest_partition( table_name, @@ -398,3 +396,26 @@ 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 == [] From 4b2f8e1fe3add0d2a02aa70bbb207b01f31a7395 Mon Sep 17 00:00:00 2001 From: Igor Khrol Date: Wed, 10 Jan 2024 16:48:16 +0200 Subject: [PATCH 4/5] Cleanup from the previous iterations --- superset/db_engine_specs/trino.py | 2 +- tests/integration_tests/databases/api_tests.py | 16 ---------------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index bec0bd8d604a5..f223777fb8942 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -59,7 +59,7 @@ def extra_table_metadata( table_name: str, schema_name: str | None, ) -> dict[str, Any]: - metadata: dict[str, Any] = {} + metadata = {} if indexes := database.get_indexes(table_name, schema_name): col_names, latest_parts = cls.latest_partition( diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index 05713c2e8d2eb..0bc1f245a1f7b 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -1340,22 +1340,6 @@ def test_get_invalid_table_table_extra_metadata(self): self.assertEqual(rv.status_code, 200) self.assertEqual(data, {}) - def test_get_invalid_table_table_extra_metadata_trino(self): - """ - Database API: Test get invalid table from table extra metadata - Trino dialect - """ - from superset.utils.database import get_or_create_db - - trino_db = get_or_create_db("trino_examples", "trino://42/sd") - uri = f"api/v1/database/{trino_db.id}/table_extra/wrong_table/null/" - self.login(username="admin") - rv = self.client.get(uri) - data = json.loads(rv.data.decode("utf-8")) - - self.assertEqual(rv.status_code, 200) - self.assertEqual(data, {}) - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_get_select_star(self): """ From d7ff579f848b239c48258aaf98339fa4d282f67d Mon Sep 17 00:00:00 2001 From: John Bodley <4567245+john-bodley@users.noreply.github.com> Date: Thu, 11 Jan 2024 11:48:44 +1300 Subject: [PATCH 5/5] Update superset/db_engine_specs/trino.py --- superset/db_engine_specs/trino.py | 1 + 1 file changed, 1 insertion(+) diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index f223777fb8942..1dc711880d729 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -407,6 +407,7 @@ def get_indexes( ) -> 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