You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "drivaspreset (via GitHub)" <gi...@apache.org> on 2026/10/15 20:49:21 UTC
[PR] test(trino): close db_engine_spec test gaps [superset]
drivaspreset opened a new pull request, #45361:
URL: https://github.com/apache/superset/pull/45361
### SUMMARY
Closes test gaps in `TrinoEngineSpec` (`superset/db_engine_specs/trino.py`). **Tests only: no production code is changed.** This replaces #44661, which had fallen far behind `master`; it is re-implemented against current `master` and addresses all four review comments from that PR (see below).
| | before | after |
|---|---|---|
| tests in `tests/unit_tests/db_engine_specs/test_trino.py` | 117 | 192 |
| `trino.py` line coverage | 90% | 98% |
(Measured with `coverage run --branch -m pytest ...` on the same commit before and after.)
`TrinoEngineSpec` extends `PrestoBaseEngineSpec`, not `PrestoEngineSpec`, so several methods it uses are only reached through the generic base classes and were never pinned at the Trino layer. New coverage:
- **Discovery:** `get_table_names` / `get_view_names` / `get_schema_names` / `get_catalog_names` (plain, empty, catalog-qualified, mixed case), and schema names containing `/`, spaces, `%`, `?`, non-ASCII round-tripping through `adjust_engine_params` / `get_schema_from_engine_params`.
- **Types:** TIME / INTERVAL mapping, unmapped IPADDRESS / UUID / HyperLogLog, ROW nested in MAP / ARRAY, deeply nested ROW expansion, `convert_dttm` for TIME / INTERVAL.
- **Exception mapping:** subclass and unmapped paths, message preservation.
- **Metadata:** `get_extra_table_metadata` for non-partitioned tables, views, partitioned-but-empty tables and tables dropped mid-inspection (through a real `Database`, so the production `NoSuchTableError` handler runs); `get_function_names`; `Database.function_names` swallowing connection errors.
- **Execution:** `cancel_query` rejecting malformed / injection-style IDs; worker-thread exception propagation in `execute_with_cursor` (before and after a query ID exists); Trino's 401 detection through `execute_with_oauth2_retry` (refresh and replay before progress, no replay after progress); the polling time limit; tracking URL; remaining small branches.
**Characterization tests.** A few tests pin what the code does today even where that behavior looks questionable, so that a future fix shows up as a deliberate test change. Each says "Characterizes current behavior" and "Invert when fixed" in its docstring. They record:
1. Driver errors that arrive wrapped in a `sqlalchemy.exc` error (as they do through a real inspector) are not mapped to Superset DBAPI exceptions, while bare driver errors are.
2. Trino defines no `custom_errors`, so every error surfaces as `GENERIC_DB_ENGINE_ERROR`.
3. `execute_with_cursor` copies Flask's `g` into the worker thread shallowly, so mutable values are shared.
4. `_expand_columns` has no recursion-depth guard.
5. A `BaseException` (e.g. `KeyboardInterrupt`) in the worker thread is not recorded, so `execute_with_cursor` returns normally.
Happy to drop any of these if reviewers would rather not pin current behavior.
**Review feedback from #44661, addressed**
- A missing-table test mocked `get_indexes` and so bypassed the production handler: it now runs through a real `Database.get_indexes`.
- Permission-denied tests asserted a bare driver exception SQLAlchemy would have wrapped: they now use `sqlalchemy.exc.ProgrammingError` with `.orig`.
- Worker-exception tests inspected `cursor.query_id` after the call: `handle_cursor` now records it at call time.
- An OAuth test's docstring described a mid-flight case it did not exercise: replaced by tests that go through `execute_with_oauth2_retry`.
**Unrelated flakiness fix in the same file.** `test_auth_custom_auth_denied` rebinds `superset.config.ALLOWED_EXTRA_AUTHENTICATIONS` and never restores it, which made `test_auth_custom_auth` fail under random test ordering (5 of 8 runs on an untouched copy of `master`'s file). It now uses `patch.dict`.
### TESTING INSTRUCTIONS
```bash
pytest tests/unit_tests/db_engine_specs/test_trino.py
```
Expected: 192 passed. The file was also run 100 times in random order, 10 times with `-n 4`, and 15 times under CPU load, with no failures. The whole `tests/unit_tests/db_engine_specs` directory was run 20 times in a row with no failures. Coverage for `trino.py`:
```bash
coverage run --branch --include='superset/db_engine_specs/trino.py' -m pytest tests/unit_tests/db_engine_specs/test_trino.py
coverage report -m
```
### ADDITIONAL INFORMATION
- [ ] Has associated issue:
- [ ] Required feature flags:
- [ ] Changes UI
- [ ] Includes DB Migration (follow approval process in [SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [ ] Introduces new feature or API
- [ ] Removes existing feature or API
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4295323424
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and marked this thread resolved** as of `435d6af`.
The test no longer treats a worker `KeyboardInterrupt` as successful execution. Worker failures are tested with `pytest.raises`, and the tests assert the original exception is re-raised.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it open.</sub>
<!-- codeant-auto-resolve-reply -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "rusackas (via GitHub)" <gi...@apache.org>.
rusackas merged PR #45361:
URL: https://github.com/apache/superset/pull/45361
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "drivaspreset (via GitHub)" <gi...@apache.org>.
drivaspreset commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4298696933
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
Review Comment:
Good catch. I checked the current failure: the RecursionError comes from `_expand_columns` itself (about 950 of the ~960 traceback frames are `_expand_columns`), not from `str(inner_type)`. But nothing in the test guaranteed that. Fixed in a0d20c636e: each nested level is now a `ROW` subclass whose `__str__` returns a fixed string, so type serialization cannot raise and the only remaining source is the expansion recursion. I verified that adding a depth guard to `_expand_columns` turns the test into a strict XPASS.
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert expanded[0] == col
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Documents current behavior; update this test if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def _run_execute_with_cursor(
+ app: Any,
+ cursor: MagicMock,
+ execute: Any,
+ g_values: dict[str, Any] | None = None,
+) -> None:
+ """
+ Run ``execute_with_cursor`` inside a request context, with
+ ``TrinoEngineSpec.execute`` (which runs in the worker thread) replaced by
+ ``execute`` as a mock side effect. ``g_values`` are set on ``g`` first.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ with (
+ app.test_request_context("/some/place/"),
+ patch.object(TrinoEngineSpec, "execute", side_effect=execute),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ for key, value in (g_values or {}).items():
+ setattr(g, key, value)
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=cursor,
+ sql="SELECT 1 FROM foo",
+ query=MagicMock(),
+ )
+
+
+def _record_query_ids_seen_by_handle_cursor(
+ mocker: MockerFixture,
+) -> list[str | None]:
+ """
+ Replace ``handle_cursor`` and record ``cursor.query_id`` at the moment it is
+ called, which is what the main thread hands to progress polling.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ observed: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed.append(cursor.query_id),
+ )
+ return observed
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=AssertionError,
+ reason="Known defect: g is copied shallowly into the worker thread.",
+)
+def test_execute_with_cursor_isolates_mutable_g_state_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A mutation the worker thread makes to a value on ``g`` should not leak back
+ into the request thread's ``g``.
+
+ Known defect, hence ``xfail(strict=True)``: the worker rebuilds ``g`` with
+ ``setattr(g, key, value)`` per attribute, a shallow copy, so a mutable value
+ (here a dict) is the same object on both sides.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="existing-query-id")
+ shared_state = {"count": 0}
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ _run_execute_with_cursor(
+ app, mock_cursor, _mutate_in_worker, g_values={"shared_state": shared_state}
+ )
+
+ assert shared_state["count"] == 0
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+
+ with pytest.raises(type(exception)) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, exception)
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with pytest.raises(RuntimeError) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, _fail_after_query_id)
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=pytest.fail.Exception,
+ reason="Known defect: the worker only records Exception subclasses.",
+)
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_reraises_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A worker that dies with a ``BaseException`` should not be reported as a
+ successful query: the error should be re-raised in the calling thread.
+
+ Known defect, hence ``xfail(strict=True)``: the worker's ``except Exception``
+ does not record ``BaseException`` subclasses, so ``execute_with_cursor``
+ returns normally. ``SystemExit`` stands in for ``KeyboardInterrupt``, which
+ pytest would treat as a request to abort the whole session.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="myQueryId")
+
+ with pytest.raises(SystemExit):
+ _run_execute_with_cursor(app, mock_cursor, SystemExit)
+
+
+def test_update_params_from_encrypted_extra_malformed_json_reraises() -> None:
+ """
+ A malformed ``encrypted_extra`` blob is logged and re-raised as
+ ``JSONDecodeError`` rather than swallowed.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = "{not valid json"
+
+ with pytest.raises(json.JSONDecodeError):
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def _oauth2_database_for_trino(
+ mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> MagicMock:
+ """
+ A mocked ``Database`` that reports OAuth2 as enabled but uses the real
+ ``TrinoEngineSpec.needs_oauth2``, so the 401 detection under test is
+ Trino's own.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock(id=1)
+ database.is_oauth2_enabled.return_value = True
+ database.db_engine_spec = TrinoEngineSpec
+ database.get_oauth2_config.return_value = oauth2_config
+ return database
+
+
+def test_expired_token_is_refreshed_and_query_retried(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised before the query has made progress (token expired) is
+ recognized by ``TrinoEngineSpec.needs_oauth2``, which makes
+ ``execute_with_oauth2_retry`` force one token refresh and replay the query.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch(
+ "superset.utils.oauth2.refresh_oauth2_token", return_value="new-token"
+ )
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=[HttpError("error 401: Unauthorized"), "rows"])
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ assert execute_with_oauth2_retry(database, operation) == "rows"
+
+ assert operation.call_count == 2
+ refresh.assert_called_once()
Review Comment:
Agreed, that was too loose. Fixed in a0d20c636e: the stored token starts as `stale-token`, the mocked refresh swaps it to `new-token`, and the operation only succeeds when it sees the new token. The test now asserts `refresh_oauth2_token(..., force=True, rejected_access_token="stale-token")`, `tokens_used == ["stale-token", "new-token"]`, and `db.session.expire(token)` (what lets the replay read the refreshed token). Verified it fails if `force=True` becomes `False` or if the expire is dropped.
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert expanded[0] == col
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Documents current behavior; update this test if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def _run_execute_with_cursor(
+ app: Any,
+ cursor: MagicMock,
+ execute: Any,
+ g_values: dict[str, Any] | None = None,
+) -> None:
+ """
+ Run ``execute_with_cursor`` inside a request context, with
+ ``TrinoEngineSpec.execute`` (which runs in the worker thread) replaced by
+ ``execute`` as a mock side effect. ``g_values`` are set on ``g`` first.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ with (
+ app.test_request_context("/some/place/"),
+ patch.object(TrinoEngineSpec, "execute", side_effect=execute),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ for key, value in (g_values or {}).items():
+ setattr(g, key, value)
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=cursor,
+ sql="SELECT 1 FROM foo",
+ query=MagicMock(),
+ )
+
+
+def _record_query_ids_seen_by_handle_cursor(
+ mocker: MockerFixture,
+) -> list[str | None]:
+ """
+ Replace ``handle_cursor`` and record ``cursor.query_id`` at the moment it is
+ called, which is what the main thread hands to progress polling.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ observed: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed.append(cursor.query_id),
+ )
+ return observed
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=AssertionError,
+ reason="Known defect: g is copied shallowly into the worker thread.",
+)
+def test_execute_with_cursor_isolates_mutable_g_state_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A mutation the worker thread makes to a value on ``g`` should not leak back
+ into the request thread's ``g``.
+
+ Known defect, hence ``xfail(strict=True)``: the worker rebuilds ``g`` with
+ ``setattr(g, key, value)`` per attribute, a shallow copy, so a mutable value
+ (here a dict) is the same object on both sides.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="existing-query-id")
+ shared_state = {"count": 0}
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ _run_execute_with_cursor(
+ app, mock_cursor, _mutate_in_worker, g_values={"shared_state": shared_state}
+ )
+
+ assert shared_state["count"] == 0
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+
+ with pytest.raises(type(exception)) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, exception)
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with pytest.raises(RuntimeError) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, _fail_after_query_id)
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=pytest.fail.Exception,
+ reason="Known defect: the worker only records Exception subclasses.",
+)
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_reraises_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A worker that dies with a ``BaseException`` should not be reported as a
+ successful query: the error should be re-raised in the calling thread.
+
+ Known defect, hence ``xfail(strict=True)``: the worker's ``except Exception``
+ does not record ``BaseException`` subclasses, so ``execute_with_cursor``
+ returns normally. ``SystemExit`` stands in for ``KeyboardInterrupt``, which
+ pytest would treat as a request to abort the whole session.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="myQueryId")
+
+ with pytest.raises(SystemExit):
+ _run_execute_with_cursor(app, mock_cursor, SystemExit)
+
+
+def test_update_params_from_encrypted_extra_malformed_json_reraises() -> None:
+ """
+ A malformed ``encrypted_extra`` blob is logged and re-raised as
+ ``JSONDecodeError`` rather than swallowed.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = "{not valid json"
+
+ with pytest.raises(json.JSONDecodeError):
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def _oauth2_database_for_trino(
+ mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> MagicMock:
+ """
+ A mocked ``Database`` that reports OAuth2 as enabled but uses the real
+ ``TrinoEngineSpec.needs_oauth2``, so the 401 detection under test is
+ Trino's own.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock(id=1)
+ database.is_oauth2_enabled.return_value = True
+ database.db_engine_spec = TrinoEngineSpec
+ database.get_oauth2_config.return_value = oauth2_config
+ return database
+
+
+def test_expired_token_is_refreshed_and_query_retried(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised before the query has made progress (token expired) is
+ recognized by ``TrinoEngineSpec.needs_oauth2``, which makes
+ ``execute_with_oauth2_retry`` force one token refresh and replay the query.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch(
+ "superset.utils.oauth2.refresh_oauth2_token", return_value="new-token"
+ )
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=[HttpError("error 401: Unauthorized"), "rows"])
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ assert execute_with_oauth2_retry(database, operation) == "rows"
+
+ assert operation.call_count == 2
+ refresh.assert_called_once()
+ database.start_oauth2_dance.assert_not_called()
+
+
+def test_expired_token_mid_query_restarts_oauth2_without_replay(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised after the query already made progress (token expired
+ mid-flight) is not replayed: ``execute_with_oauth2_retry`` skips the refresh
+ and starts the OAuth2 dance so the user re-authenticates. This is the
+ behavior that makes long-running queries lose in-flight progress.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch("superset.utils.oauth2.refresh_oauth2_token")
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=HttpError("error 401: Unauthorized"))
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ with pytest.raises(HttpError):
+ execute_with_oauth2_retry(database, operation, can_retry=lambda: False)
+
+ operation.assert_called_once_with()
+ refresh.assert_not_called()
+ database.start_oauth2_dance.assert_called_once_with()
+
+
+# ---------------------------------------------------------------------------
+# Remaining small branches of trino.py: view/partition metadata, tracking URL,
+# cost-estimate gate, encrypted-extra no-ops, and progress-polling timeout.
+# ---------------------------------------------------------------------------
+
+
+def test_get_extra_table_metadata_includes_view_definition(
+ mocker: MockerFixture,
+) -> None:
+ """
+ For a view, the definition comes from the inspector and is returned under
+ the ``view`` key.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.has_view = Mock(return_value=True)
+ inspector = db_mock.get_inspector.return_value.__enter__.return_value
+ inspector.get_view_definition.return_value = "SELECT 1"
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("my_view", "test_schema")
+ )
+
+ assert result == {"view": "SELECT 1"}
+ inspector.get_view_definition.assert_called_once_with("my_view", "test_schema")
+
+
+def test_get_extra_table_metadata_partitioned_table_without_data(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A partitioned table with no partitions yet reports its partition columns,
+ with every latest value ``None``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(
+ return_value=[{"column_names": ["ds", "hour"], "name": "partition"}]
+ )
+ db_mock.has_view = Mock(return_value=False)
+ mocker.patch.object(
+ TrinoEngineSpec, "latest_partition", return_value=(["ds", "hour"], None)
+ )
+ mocker.patch.object(TrinoEngineSpec, "_partition_query", return_value="SELECT 1")
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result["partitions"]["cols"] == ["ds", "hour"]
+ assert result["partitions"]["latest"] == {"ds": None, "hour": None}
+
+
+def test_latest_sub_partition_reads_indexes_from_database(
+ mocker: MockerFixture,
+) -> None:
+ """
+ When no indexes are passed, ``latest_sub_partition`` fetches them from the
+ database and hands the (Iceberg-filtered) indexes to the base implementation.
+ """
+ from superset.db_engine_specs.presto import PrestoBaseEngineSpec
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ indexes = [{"name": "partition", "column_names": ["ds"]}]
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=indexes)
+ base_mock = mocker.patch.object(
+ PrestoBaseEngineSpec, "latest_sub_partition", return_value="2019-01-01"
+ )
+ table = Table("test_table", "test_schema")
+
+ assert TrinoEngineSpec.latest_sub_partition(db_mock, table, ds="x") == "2019-01-01"
+
+ db_mock.get_indexes.assert_called_once_with(table)
+ base_mock.assert_called_once_with(db_mock, table, indexes=indexes, ds="x")
+
+
+def test_impersonate_user_without_username_returns_inputs_unchanged() -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://localhost:8080/system")
+ engine_kwargs: dict[str, Any] = {}
+
+ new_url, new_kwargs = TrinoEngineSpec.impersonate_user(
+ Mock(), None, "token", url, engine_kwargs
+ )
+
+ assert new_url is url
+ assert new_kwargs == {}
+
+
+def test_get_allow_cost_estimate_is_always_true() -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_allow_cost_estimate({}) is True
+ assert TrinoEngineSpec.get_allow_cost_estimate({"version": "0.1"}) is True
+
+
+def test_get_tracking_url_prefers_info_uri(mocker: MockerFixture) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock()
+ cursor.info_uri = "http://trino:8080/ui/query.html?abc"
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) == cursor.info_uri
+
+
+def test_get_tracking_url_builds_url_when_info_uri_is_unavailable(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Older drivers expose no ``info_uri``; the URL is rebuilt from the
+ connection's scheme/host/port and the cursor's query ID.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock(spec=["connection", "_query"])
+ cursor.connection.http_scheme = "https"
+ cursor.connection.host = "trino.example.com"
+ cursor.connection.port = 8443
+ cursor._query.query_id = "20240101_000000_00001_abcde"
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) == (
+ "https://trino.example.com:8443/ui/query.html?20240101_000000_00001_abcde"
+ )
+
+
+def test_get_tracking_url_is_none_without_connection_details(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock(spec=[])
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) is None
+
+
+@pytest.mark.parametrize(
+ "encrypted_extra",
+ [
+ None,
+ "",
+ # valid JSON but no auth_method: nothing to configure
+ '{"auth_params": {"username": "u", "password": "p"}}',
+ ],
+)
+def test_update_params_from_encrypted_extra_without_auth_method_is_noop(
+ encrypted_extra: str | None,
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = encrypted_extra
+ params: dict[str, Any] = {}
+
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, params)
+
+ assert params == {}
+
+
+@patch("superset.db_engine_specs.presto.PrestoBaseEngineSpec.handle_cursor")
+@patch("superset.db_engine_specs.trino.TrinoEngineSpec.cancel_query")
+@patch("superset.db_engine_specs.trino.db")
+@patch("superset.db_engine_specs.trino.app")
+def test_handle_cursor_stops_polling_after_time_limit(
+ mock_app: Mock,
+ mock_db: Mock,
+ mock_cancel_query: Mock,
+ mock_presto_handle_cursor: Mock,
+ mocker: MockerFixture,
+) -> None:
+ """
+ Progress polling gives up once ``SQLLAB_ASYNC_TIME_LIMIT_SEC`` has elapsed
+ instead of looping forever on a query that never reaches a terminal state.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ mock_app.config = {
+ "DB_POLL_INTERVAL_SECONDS": {"trino": 0},
+ "SQLLAB_ASYNC_TIME_LIMIT_SEC": -1,
Review Comment:
Yes, `-1` only proved the early exit. Fixed in a0d20c636e: limit is 10 and the module clock is controlled (`time.time` returns 0 at start, 5 on the first check, 11 on the second). The test asserts progress was updated to 50 on the in-limit pass, exactly one poll sleep, and no cancel. Verified it fails if `start_time` is reset on every pass (the clock runs out instead of the loop stopping).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codecov[bot] (via GitHub)" <gi...@apache.org>.
codecov[bot] commented on PR #45361:
URL: https://github.com/apache/superset/pull/45361#issuecomment-6203100722
## [Codecov](https://app.codecov.io/gh/apache/superset/pull/45361?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 64.91%. Comparing base ([`e0e728e`](https://app.codecov.io/gh/apache/superset/commit/e0e728ef7644451a8b7162a7130175fdadef1571?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`b211fdd`](https://app.codecov.io/gh/apache/superset/commit/b211fdde7a68e1d3b768699906539bf7a3e1bfcc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 3 commits behind head on master.
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #45361 +/- ##
==========================================
+ Coverage 56.21% 64.91% +8.70%
==========================================
Files 3045 3045
Lines 199194 199199 +5
Branches 46455 46457 +2
==========================================
+ Hits 111984 129319 +17335
+ Misses 86143 66967 -19176
- Partials 1067 2913 +1846
```
| [Flag](https://app.codecov.io/gh/apache/superset/pull/45361/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [hive](https://app.codecov.io/gh/apache/superset/pull/45361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `35.00% <ø> (-0.01%)` | :arrow_down: |
| [presto](https://app.codecov.io/gh/apache/superset/pull/45361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `36.76% <ø> (-0.01%)` | :arrow_down: |
| [python](https://app.codecov.io/gh/apache/superset/pull/45361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `52.85% <ø> (+16.04%)` | :arrow_up: |
| [sqlite](https://app.codecov.io/gh/apache/superset/pull/45361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `52.73% <ø> (?)` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/superset/pull/45361?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4298705718
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and marked this thread resolved** as of `a0d20c6`.
The worker now mutates `g.shared_state`, and the test is marked `xfail(strict=True, raises=AssertionError)` while asserting the caller’s state remains unchanged. It no longer treats a worker `KeyboardInterrupt` as successful execution.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it open.</sub>
<!-- codeant-auto-resolve-reply -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4288226395
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
**Suggestion:** This assertion treats a worker `KeyboardInterrupt` as successful execution, so silent query failure becomes the test's required behavior.
**Assessment:** 🟠 `Major` · 🔁 `Occurrence: Rarely` · 🏷️ `Incomplete implementation`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=017805f625fb443e8966d6532ccabdff&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=017805f625fb443e8966d6532ccabdff&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** tests/unit_tests/db_engine_specs/test_trino.py
**Line:** 2679:2679
**Comment:**
*Incomplete Implementation: This assertion treats a worker `KeyboardInterrupt` as successful execution, so silent query failure becomes the test's required behavior.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45361&comment_hash=c0b76fc24b7b40ad27b1b215d7a969f5d44001dfab725363d72e8c740721e4a4&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45361&comment_hash=c0b76fc24b7b40ad27b1b215d7a969f5d44001dfab725363d72e8c740721e4a4&reaction=dislike'>👎</a>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "sadpandajoe (via GitHub)" <gi...@apache.org>.
sadpandajoe commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4297894176
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
Review Comment:
`_expand_columns` calls `str(inner_type)` (for `get_column_spec`) on the full nested ROW before it recurses, so with a ROW deeper than the recursion limit the `RecursionError` can come from type serialization rather than from the missing expansion guard. If so, a depth guard in `_expand_columns` would still leave this test raising `RecursionError`, and the strict xfail would never flip. Could the test make sure the failure is attributable to the expansion recursion, e.g. by patching `get_column_spec` or asserting where the error is raised?
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert expanded[0] == col
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Documents current behavior; update this test if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def _run_execute_with_cursor(
+ app: Any,
+ cursor: MagicMock,
+ execute: Any,
+ g_values: dict[str, Any] | None = None,
+) -> None:
+ """
+ Run ``execute_with_cursor`` inside a request context, with
+ ``TrinoEngineSpec.execute`` (which runs in the worker thread) replaced by
+ ``execute`` as a mock side effect. ``g_values`` are set on ``g`` first.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ with (
+ app.test_request_context("/some/place/"),
+ patch.object(TrinoEngineSpec, "execute", side_effect=execute),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ for key, value in (g_values or {}).items():
+ setattr(g, key, value)
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=cursor,
+ sql="SELECT 1 FROM foo",
+ query=MagicMock(),
+ )
+
+
+def _record_query_ids_seen_by_handle_cursor(
+ mocker: MockerFixture,
+) -> list[str | None]:
+ """
+ Replace ``handle_cursor`` and record ``cursor.query_id`` at the moment it is
+ called, which is what the main thread hands to progress polling.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ observed: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed.append(cursor.query_id),
+ )
+ return observed
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=AssertionError,
+ reason="Known defect: g is copied shallowly into the worker thread.",
+)
+def test_execute_with_cursor_isolates_mutable_g_state_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A mutation the worker thread makes to a value on ``g`` should not leak back
+ into the request thread's ``g``.
+
+ Known defect, hence ``xfail(strict=True)``: the worker rebuilds ``g`` with
+ ``setattr(g, key, value)`` per attribute, a shallow copy, so a mutable value
+ (here a dict) is the same object on both sides.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="existing-query-id")
+ shared_state = {"count": 0}
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ _run_execute_with_cursor(
+ app, mock_cursor, _mutate_in_worker, g_values={"shared_state": shared_state}
+ )
+
+ assert shared_state["count"] == 0
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+
+ with pytest.raises(type(exception)) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, exception)
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with pytest.raises(RuntimeError) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, _fail_after_query_id)
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=pytest.fail.Exception,
+ reason="Known defect: the worker only records Exception subclasses.",
+)
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_reraises_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A worker that dies with a ``BaseException`` should not be reported as a
+ successful query: the error should be re-raised in the calling thread.
+
+ Known defect, hence ``xfail(strict=True)``: the worker's ``except Exception``
+ does not record ``BaseException`` subclasses, so ``execute_with_cursor``
+ returns normally. ``SystemExit`` stands in for ``KeyboardInterrupt``, which
+ pytest would treat as a request to abort the whole session.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="myQueryId")
+
+ with pytest.raises(SystemExit):
+ _run_execute_with_cursor(app, mock_cursor, SystemExit)
+
+
+def test_update_params_from_encrypted_extra_malformed_json_reraises() -> None:
+ """
+ A malformed ``encrypted_extra`` blob is logged and re-raised as
+ ``JSONDecodeError`` rather than swallowed.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = "{not valid json"
+
+ with pytest.raises(json.JSONDecodeError):
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def _oauth2_database_for_trino(
+ mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> MagicMock:
+ """
+ A mocked ``Database`` that reports OAuth2 as enabled but uses the real
+ ``TrinoEngineSpec.needs_oauth2``, so the 401 detection under test is
+ Trino's own.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock(id=1)
+ database.is_oauth2_enabled.return_value = True
+ database.db_engine_spec = TrinoEngineSpec
+ database.get_oauth2_config.return_value = oauth2_config
+ return database
+
+
+def test_expired_token_is_refreshed_and_query_retried(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised before the query has made progress (token expired) is
+ recognized by ``TrinoEngineSpec.needs_oauth2``, which makes
+ ``execute_with_oauth2_retry`` force one token refresh and replay the query.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch(
+ "superset.utils.oauth2.refresh_oauth2_token", return_value="new-token"
+ )
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=[HttpError("error 401: Unauthorized"), "rows"])
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ assert execute_with_oauth2_retry(database, operation) == "rows"
+
+ assert operation.call_count == 2
+ refresh.assert_called_once()
Review Comment:
`refresh.assert_called_once()` passes whether or not the refresh is forced, and the mocked operation returns `"rows"` on the second call regardless of which token it is given. If the retry path regressed to an ordinary refresh that reuses the rejected-but-unexpired token, this test would still pass. Should this assert the refresh was called with `force=True` (and that the replay saw the new token)?
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -2020,3 +2040,986 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+def _show_catalogs_connection(inspector: MagicMock) -> MagicMock:
+ """The connection ``get_catalog_names`` opens to run ``SHOW CATALOGS``."""
+ return inspector.engine.connect.return_value.__enter__.return_value
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Documents current behavior; the intended behavior is not decided yet.
+ Update this test if the mapping learns to look at ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = _show_catalogs_connection(inspector)
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=RecursionError,
+ reason="Known defect: _expand_columns has no recursion-depth guard.",
+)
+def test_expand_columns_survives_rows_nested_beyond_recursion_limit() -> None:
+ """
+ Expanding a ROW nested deeper than the interpreter's recursion limit should
+ not crash column inspection for the whole table.
+
+ Known defect, hence ``xfail(strict=True)``: ``_expand_columns`` recurses
+ once per ROW level with no maximum depth and raises ``RecursionError``. Once
+ a guard lands this test passes, strict mode turns that into a failure, and
+ the marker should be removed. The type is built programmatically because
+ ``parse_sqltype`` is itself recursive and would fail first on a string this
+ deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert expanded[0] == col
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Documents current behavior; update this test if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def _run_execute_with_cursor(
+ app: Any,
+ cursor: MagicMock,
+ execute: Any,
+ g_values: dict[str, Any] | None = None,
+) -> None:
+ """
+ Run ``execute_with_cursor`` inside a request context, with
+ ``TrinoEngineSpec.execute`` (which runs in the worker thread) replaced by
+ ``execute`` as a mock side effect. ``g_values`` are set on ``g`` first.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ with (
+ app.test_request_context("/some/place/"),
+ patch.object(TrinoEngineSpec, "execute", side_effect=execute),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ for key, value in (g_values or {}).items():
+ setattr(g, key, value)
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=cursor,
+ sql="SELECT 1 FROM foo",
+ query=MagicMock(),
+ )
+
+
+def _record_query_ids_seen_by_handle_cursor(
+ mocker: MockerFixture,
+) -> list[str | None]:
+ """
+ Replace ``handle_cursor`` and record ``cursor.query_id`` at the moment it is
+ called, which is what the main thread hands to progress polling.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ observed: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed.append(cursor.query_id),
+ )
+ return observed
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=AssertionError,
+ reason="Known defect: g is copied shallowly into the worker thread.",
+)
+def test_execute_with_cursor_isolates_mutable_g_state_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A mutation the worker thread makes to a value on ``g`` should not leak back
+ into the request thread's ``g``.
+
+ Known defect, hence ``xfail(strict=True)``: the worker rebuilds ``g`` with
+ ``setattr(g, key, value)`` per attribute, a shallow copy, so a mutable value
+ (here a dict) is the same object on both sides.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="existing-query-id")
+ shared_state = {"count": 0}
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ _run_execute_with_cursor(
+ app, mock_cursor, _mutate_in_worker, g_values={"shared_state": shared_state}
+ )
+
+ assert shared_state["count"] == 0
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+
+ with pytest.raises(type(exception)) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, exception)
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ observed_query_ids = _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id=None)
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with pytest.raises(RuntimeError) as exc_info:
+ _run_execute_with_cursor(app, mock_cursor, _fail_after_query_id)
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=pytest.fail.Exception,
+ reason="Known defect: the worker only records Exception subclasses.",
+)
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_reraises_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ A worker that dies with a ``BaseException`` should not be reported as a
+ successful query: the error should be re-raised in the calling thread.
+
+ Known defect, hence ``xfail(strict=True)``: the worker's ``except Exception``
+ does not record ``BaseException`` subclasses, so ``execute_with_cursor``
+ returns normally. ``SystemExit`` stands in for ``KeyboardInterrupt``, which
+ pytest would treat as a request to abort the whole session.
+ """
+ _record_query_ids_seen_by_handle_cursor(mocker)
+ mock_cursor = mocker.MagicMock(query_id="myQueryId")
+
+ with pytest.raises(SystemExit):
+ _run_execute_with_cursor(app, mock_cursor, SystemExit)
+
+
+def test_update_params_from_encrypted_extra_malformed_json_reraises() -> None:
+ """
+ A malformed ``encrypted_extra`` blob is logged and re-raised as
+ ``JSONDecodeError`` rather than swallowed.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = "{not valid json"
+
+ with pytest.raises(json.JSONDecodeError):
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, {})
+
+
+def _oauth2_database_for_trino(
+ mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> MagicMock:
+ """
+ A mocked ``Database`` that reports OAuth2 as enabled but uses the real
+ ``TrinoEngineSpec.needs_oauth2``, so the 401 detection under test is
+ Trino's own.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock(id=1)
+ database.is_oauth2_enabled.return_value = True
+ database.db_engine_spec = TrinoEngineSpec
+ database.get_oauth2_config.return_value = oauth2_config
+ return database
+
+
+def test_expired_token_is_refreshed_and_query_retried(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised before the query has made progress (token expired) is
+ recognized by ``TrinoEngineSpec.needs_oauth2``, which makes
+ ``execute_with_oauth2_retry`` force one token refresh and replay the query.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch(
+ "superset.utils.oauth2.refresh_oauth2_token", return_value="new-token"
+ )
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=[HttpError("error 401: Unauthorized"), "rows"])
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ assert execute_with_oauth2_retry(database, operation) == "rows"
+
+ assert operation.call_count == 2
+ refresh.assert_called_once()
+ database.start_oauth2_dance.assert_not_called()
+
+
+def test_expired_token_mid_query_restarts_oauth2_without_replay(
+ app, mocker: MockerFixture, oauth2_config: OAuth2ClientConfig
+) -> None:
+ """
+ A Trino 401 raised after the query already made progress (token expired
+ mid-flight) is not replayed: ``execute_with_oauth2_retry`` skips the refresh
+ and starts the OAuth2 dance so the user re-authenticates. This is the
+ behavior that makes long-running queries lose in-flight progress.
+ """
+ from superset.utils.oauth2 import execute_with_oauth2_retry
+
+ database = _oauth2_database_for_trino(mocker, oauth2_config)
+ refresh = mocker.patch("superset.utils.oauth2.refresh_oauth2_token")
+ mocker.patch("superset.utils.oauth2.db")
+ operation = mocker.Mock(side_effect=HttpError("error 401: Unauthorized"))
+
+ with app.test_request_context("/"):
+ g.user = mocker.Mock(id=2)
+ with pytest.raises(HttpError):
+ execute_with_oauth2_retry(database, operation, can_retry=lambda: False)
+
+ operation.assert_called_once_with()
+ refresh.assert_not_called()
+ database.start_oauth2_dance.assert_called_once_with()
+
+
+# ---------------------------------------------------------------------------
+# Remaining small branches of trino.py: view/partition metadata, tracking URL,
+# cost-estimate gate, encrypted-extra no-ops, and progress-polling timeout.
+# ---------------------------------------------------------------------------
+
+
+def test_get_extra_table_metadata_includes_view_definition(
+ mocker: MockerFixture,
+) -> None:
+ """
+ For a view, the definition comes from the inspector and is returned under
+ the ``view`` key.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.has_view = Mock(return_value=True)
+ inspector = db_mock.get_inspector.return_value.__enter__.return_value
+ inspector.get_view_definition.return_value = "SELECT 1"
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("my_view", "test_schema")
+ )
+
+ assert result == {"view": "SELECT 1"}
+ inspector.get_view_definition.assert_called_once_with("my_view", "test_schema")
+
+
+def test_get_extra_table_metadata_partitioned_table_without_data(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A partitioned table with no partitions yet reports its partition columns,
+ with every latest value ``None``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(
+ return_value=[{"column_names": ["ds", "hour"], "name": "partition"}]
+ )
+ db_mock.has_view = Mock(return_value=False)
+ mocker.patch.object(
+ TrinoEngineSpec, "latest_partition", return_value=(["ds", "hour"], None)
+ )
+ mocker.patch.object(TrinoEngineSpec, "_partition_query", return_value="SELECT 1")
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result["partitions"]["cols"] == ["ds", "hour"]
+ assert result["partitions"]["latest"] == {"ds": None, "hour": None}
+
+
+def test_latest_sub_partition_reads_indexes_from_database(
+ mocker: MockerFixture,
+) -> None:
+ """
+ When no indexes are passed, ``latest_sub_partition`` fetches them from the
+ database and hands the (Iceberg-filtered) indexes to the base implementation.
+ """
+ from superset.db_engine_specs.presto import PrestoBaseEngineSpec
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ indexes = [{"name": "partition", "column_names": ["ds"]}]
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=indexes)
+ base_mock = mocker.patch.object(
+ PrestoBaseEngineSpec, "latest_sub_partition", return_value="2019-01-01"
+ )
+ table = Table("test_table", "test_schema")
+
+ assert TrinoEngineSpec.latest_sub_partition(db_mock, table, ds="x") == "2019-01-01"
+
+ db_mock.get_indexes.assert_called_once_with(table)
+ base_mock.assert_called_once_with(db_mock, table, indexes=indexes, ds="x")
+
+
+def test_impersonate_user_without_username_returns_inputs_unchanged() -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://localhost:8080/system")
+ engine_kwargs: dict[str, Any] = {}
+
+ new_url, new_kwargs = TrinoEngineSpec.impersonate_user(
+ Mock(), None, "token", url, engine_kwargs
+ )
+
+ assert new_url is url
+ assert new_kwargs == {}
+
+
+def test_get_allow_cost_estimate_is_always_true() -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_allow_cost_estimate({}) is True
+ assert TrinoEngineSpec.get_allow_cost_estimate({"version": "0.1"}) is True
+
+
+def test_get_tracking_url_prefers_info_uri(mocker: MockerFixture) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock()
+ cursor.info_uri = "http://trino:8080/ui/query.html?abc"
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) == cursor.info_uri
+
+
+def test_get_tracking_url_builds_url_when_info_uri_is_unavailable(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Older drivers expose no ``info_uri``; the URL is rebuilt from the
+ connection's scheme/host/port and the cursor's query ID.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock(spec=["connection", "_query"])
+ cursor.connection.http_scheme = "https"
+ cursor.connection.host = "trino.example.com"
+ cursor.connection.port = 8443
+ cursor._query.query_id = "20240101_000000_00001_abcde"
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) == (
+ "https://trino.example.com:8443/ui/query.html?20240101_000000_00001_abcde"
+ )
+
+
+def test_get_tracking_url_is_none_without_connection_details(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ cursor = mocker.MagicMock(spec=[])
+
+ assert TrinoEngineSpec.get_tracking_url(cursor) is None
+
+
+@pytest.mark.parametrize(
+ "encrypted_extra",
+ [
+ None,
+ "",
+ # valid JSON but no auth_method: nothing to configure
+ '{"auth_params": {"username": "u", "password": "p"}}',
+ ],
+)
+def test_update_params_from_encrypted_extra_without_auth_method_is_noop(
+ encrypted_extra: str | None,
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = Mock()
+ database.encrypted_extra = encrypted_extra
+ params: dict[str, Any] = {}
+
+ TrinoEngineSpec.update_params_from_encrypted_extra(database, params)
+
+ assert params == {}
+
+
+@patch("superset.db_engine_specs.presto.PrestoBaseEngineSpec.handle_cursor")
+@patch("superset.db_engine_specs.trino.TrinoEngineSpec.cancel_query")
+@patch("superset.db_engine_specs.trino.db")
+@patch("superset.db_engine_specs.trino.app")
+def test_handle_cursor_stops_polling_after_time_limit(
+ mock_app: Mock,
+ mock_db: Mock,
+ mock_cancel_query: Mock,
+ mock_presto_handle_cursor: Mock,
+ mocker: MockerFixture,
+) -> None:
+ """
+ Progress polling gives up once ``SQLLAB_ASYNC_TIME_LIMIT_SEC`` has elapsed
+ instead of looping forever on a query that never reaches a terminal state.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ mock_app.config = {
+ "DB_POLL_INTERVAL_SECONDS": {"trino": 0},
+ "SQLLAB_ASYNC_TIME_LIMIT_SEC": -1,
Review Comment:
With `SQLLAB_ASYNC_TIME_LIMIT_SEC` of `-1` the loop exits before its first iteration, so `query.progress is None` would also hold if `start_time` were reset on every pass and a positive deadline could never expire. Would a positive limit with a controlled clock (one poll, then past the deadline) pin the deadline-crossing behavior, asserting progress was updated before polling stopped?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "netlify[bot] (via GitHub)" <gi...@apache.org>.
netlify[bot] commented on PR #45361:
URL: https://github.com/apache/superset/pull/45361#issuecomment-6202712038
### <span aria-hidden="true">✅</span> Deploy Preview for *superset-docs-preview* ready!
| Name | Link |
|:-:|------------------------|
|<span aria-hidden="true">🔨</span> Latest commit | b211fdde7a68e1d3b768699906539bf7a3e1bfcc |
|<span aria-hidden="true">🔍</span> Latest deploy log | https://app.netlify.com/projects/superset-docs-preview/deploys/6ad145cd60716e72a677b721 |
|<span aria-hidden="true">😎</span> Deploy Preview | [https://deploy-preview-45361--superset-docs-preview.netlify.app](https://deploy-preview-45361--superset-docs-preview.netlify.app) |
|<span aria-hidden="true">📱</span> Preview on mobile | <details><summary> Toggle QR Code... </summary><br /><br /><br /><br />_Use your smartphone camera to open QR code link._</details> |
|<span aria-hidden="true">🤖</span> Make changes | [Run an agent on this branch](https://app.netlify.com/projects/superset-docs-preview/agent-runs#q1YqKMrPLShRslJS0lFKKkrMS84AsotTUkt0S4oy8_J1S1KLS3TTEwuKgfIpqQU5-ZWeKUAVZokphiamySlmBuaGZqnmRolm5uZJ5kaGSrUA) |
---
<!-- [superset-docs-preview Preview](https://deploy-preview-45361--superset-docs-preview.netlify.app) -->
_To edit notification comments on pull requests, go to your [Netlify project configuration](https://app.netlify.com/projects/superset-docs-preview/configuration/notifications#deploy-notifications)._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on PR #45361:
URL: https://github.com/apache/superset/pull/45361#issuecomment-6202178825
## CodeAnt PR Risk: Medium Risk
- The PR needs attention before merging: a new test requires a worker KeyboardInterrupt to return normally, treating failed query execution as success.
- Other characterization tests assert known defects, including unbounded ROW recursion and shared mutable state across threads, so fixing those behaviors would break the tests.
**Assessed commit:** `98277ea304f2`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "drivaspreset (via GitHub)" <gi...@apache.org>.
drivaspreset commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4288476704
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
Agreed, the earlier version required the buggy behavior. Fixed in 3a: this test (and the two other known-defect tests: `_expand_columns` recursion depth and the shallow `g` copy) now asserts the intended behavior and is marked `xfail(strict=True)`, following `tests/unit_tests/charts/test_non_additive_totals.py`. A fix makes it XPASS, which strict mode reports as a failure so the marker gets removed.
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
Agreed, the earlier version required the buggy behavior. Fixed in b211fdde7a: this test (and the two other known-defect tests: `_expand_columns` recursion depth and the shallow `g` copy) now asserts the intended behavior and is marked `xfail(strict=True)`, following `tests/unit_tests/charts/test_non_additive_totals.py`. A fix makes it XPASS, which strict mode reports as a failure so the marker gets removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org
Re: [PR] test(trino): close db_engine_spec test gaps [superset]
Posted by "codeant-ai-for-open-source[bot] (via GitHub)" <gi...@apache.org>.
codeant-ai-for-open-source[bot] commented on code in PR #45361:
URL: https://github.com/apache/superset/pull/45361#discussion_r4317380933
##########
tests/unit_tests/db_engine_specs/test_trino.py:
##########
@@ -1968,3 +1988,978 @@ def test_get_cancel_query_id_reads_the_running_query() -> None:
== "20261001_000000_00001_abcde"
)
assert TrinoEngineSpec.has_query_id_during_execute is True
+
+
+# ---------------------------------------------------------------------------
+# Schema / catalog / table / view discovery.
+#
+# TrinoEngineSpec(PrestoBaseEngineSpec) does not inherit PrestoEngineSpec, so
+# get_table_names/get_view_names/get_schema_names resolve to BaseEngineSpec's
+# generic inspector-based implementation, and get_catalog_names/
+# get_function_names to PrestoBaseEngineSpec. These tests pin that resolution
+# at the Trino layer (the Presto unit tests only exercise PrestoEngineSpec).
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def trino_database(mocker: MockerFixture) -> tuple[Any, MagicMock]:
+ """
+ A real ``Database`` on the Trino engine spec whose inspector is mocked.
+
+ Returns the database and the inspector handed out by ``get_inspector``, so
+ tests can exercise ``Database`` -> ``TrinoEngineSpec`` call chains (including
+ exception mapping) without a live Trino server.
+ """
+ from superset.models.core import Database
+
+ inspector = mocker.MagicMock()
+ mocker.patch.object(
+ Database, "get_inspector"
+ ).return_value.__enter__.return_value = inspector
+ database = Database(
+ database_name="trino_db",
+ sqlalchemy_uri="trino://user@localhost:8080/catalog",
+ )
+ return database, inspector
+
+
+def _driver_error_wrapped_by_sqlalchemy(orig: Exception) -> ProgrammingError:
+ """
+ Wrap a driver exception the way SQLAlchemy does when it executes a
+ statement: a ``sqlalchemy.exc`` error with the driver exception as ``.orig``.
+ """
+ return ProgrammingError("SELECT 1", {}, orig)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+@pytest.mark.parametrize(
+ "inspector_result,schema,expected",
+ [
+ (["orders", "customers"], "my_schema", {"orders", "customers"}),
+ ([], "empty_schema", set()),
+ # a catalog-qualified schema prefix is stripped from matching names only
+ (
+ ["my_catalog.my_schema.orders", "unrelated"],
+ "my_catalog.my_schema",
+ {"orders", "unrelated"},
+ ),
+ # mixed-case names are returned verbatim
+ (["MixedCase", "lower"], "my_schema", {"MixedCase", "lower"}),
+ ],
+)
+def test_get_table_and_view_names(
+ mocker: MockerFixture,
+ method: str,
+ inspector_result: list[str],
+ schema: str,
+ expected: set[str],
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).return_value = inspector_result
+
+ assert getattr(TrinoEngineSpec, method)(Mock(), inspector, schema) == expected
+ getattr(inspector, method).assert_called_once_with(schema)
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_map_driver_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ A driver error raised while listing a schema (e.g. the schema does not
+ exist) is mapped to a Superset DBAPI exception rather than leaked raw.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = TrinoUserError(
+ {"message": "Schema 'nope' does not exist"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="does not exist"):
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+
+@pytest.mark.parametrize("method", ["get_table_names", "get_view_names"])
+def test_get_table_and_view_names_do_not_map_sqlalchemy_wrapped_error(
+ mocker: MockerFixture, method: str
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ The exception mapping keys on the driver exception type, but when the error
+ surfaces through a real SQLAlchemy inspector it arrives wrapped in a
+ ``sqlalchemy.exc`` error (driver error preserved as ``.orig``). The wrapper
+ is not in the mapping, so it propagates unmapped.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ getattr(inspector, method).side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Schema 'nope' does not exist"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ getattr(TrinoEngineSpec, method)(Mock(), inspector, "nope")
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "inspector_result,expected",
+ [
+ (["information_schema", "default"], {"information_schema", "default"}),
+ ([], set()),
+ ],
+)
+def test_get_schema_names(
+ mocker: MockerFixture, inspector_result: list[str], expected: set[str]
+) -> None:
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ inspector.get_schema_names.return_value = inspector_result
+
+ assert TrinoEngineSpec.get_schema_names(inspector) == expected
+
+
+@pytest.mark.parametrize(
+ "rows,expected",
+ [
+ ([("system",), ("tpch",), ("memory",)], {"system", "tpch", "memory"}),
+ ([], set()),
+ ],
+)
+def test_get_catalog_names(
+ mocker: MockerFixture, rows: list[tuple[str]], expected: set[str]
+) -> None:
+ """
+ get_catalog_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_catalog_names_lists_catalogs at the Trino layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.return_value = rows
+
+ assert TrinoEngineSpec.get_catalog_names(Mock(), inspector) == expected
+ assert str(conn.execute.call_args[0][0]) == "SHOW CATALOGS"
+
+
+def test_get_all_schema_names_maps_bare_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Listing schemas without access to the catalog surfaces as a Superset
+ DBAPI exception when the driver error reaches ``Database`` unwrapped.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = TrinoUserError(
+ {"message": "Access Denied: Cannot access catalog system"}
+ )
+
+ with pytest.raises(SupersetDBAPIProgrammingError, match="Access Denied"):
+ database.get_all_schema_names()
+
+
+def test_get_all_schema_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ Through a real inspector, a permission failure arrives wrapped in a
+ ``sqlalchemy.exc.ProgrammingError`` (driver error as ``.orig``). The
+ mapping does not recognize the wrapper, so the raw SQLAlchemy error is
+ raised instead of a Superset DBAPI exception.
+ """
+ database, inspector = trino_database
+ inspector.get_schema_names.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot access catalog system"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ database.get_all_schema_names()
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+def test_get_catalog_names_does_not_map_sqlalchemy_wrapped_permission_denied(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Characterizes current behavior. Invert if the mapping learns to look at
+ ``.orig``.
+
+ ``get_catalog_names`` runs ``SHOW CATALOGS`` itself and does not map
+ failures at all, so a SQLAlchemy-wrapped permission error propagates with
+ the driver error preserved as ``.orig``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ inspector = mocker.MagicMock()
+ conn = inspector.engine.connect.return_value.__enter__.return_value
+ conn.execute.side_effect = _driver_error_wrapped_by_sqlalchemy(
+ TrinoUserError({"message": "Access Denied: Cannot execute query"})
+ )
+
+ with pytest.raises(ProgrammingError) as exc_info:
+ TrinoEngineSpec.get_catalog_names(Mock(), inspector)
+
+ assert isinstance(exc_info.value.orig, TrinoUserError)
+
+
+@pytest.mark.parametrize(
+ "schema",
+ ["my schema", "a/b", "100%", "q?x=1&y=2", "ünïcödé", "with.dot"],
+)
+def test_schema_name_round_trips_through_engine_params(schema: str) -> None:
+ """
+ Schema names containing URL-significant characters survive being written
+ into the SQLAlchemy URI by ``adjust_engine_params`` and read back by
+ ``get_schema_from_engine_params``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ url = make_url("trino://user:pass@localhost:8080/system")
+
+ adjusted, _ = TrinoEngineSpec.adjust_engine_params(url, {}, schema=schema)
+
+ assert TrinoEngineSpec.get_schema_from_engine_params(adjusted, {}) == schema
+
+
+# ---------------------------------------------------------------------------
+# Type handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("native_type", ["IPADDRESS", "UUID", "HyperLogLog"])
+def test_get_column_spec_unmapped_trino_types(native_type: str) -> None:
+ """
+ Trino-only opaque types with no entry in ``column_type_mappings`` yield
+ ``None`` rather than raising, like any other unrecognized native type.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ assert TrinoEngineSpec.get_column_spec(native_type) is None
+
+
+@pytest.mark.parametrize(
+ "type_str",
+ [
+ "map(varchar, row(a varchar, b date))",
+ "array(row(a varchar, b date))",
+ ],
+)
+def test_expand_columns_does_not_expand_row_nested_in_map_or_array(
+ type_str: str,
+) -> None:
+ """
+ ``_expand_columns`` only descends into ROW types directly; a ROW nested
+ under a MAP or ARRAY is left untouched, as documented on the method.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col = ResultSetColumnType(
+ name="field1",
+ column_name="field1",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ assert TrinoEngineSpec._expand_columns(col) == [col]
+
+
+def test_expand_columns_expands_deeply_nested_row() -> None:
+ """
+ A deeply nested ROW expands to one dotted column per nesting level.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ depth = 20
+ type_str = "varchar"
+ for _ in range(depth):
+ type_str = f"row(a {type_str})"
+ col = ResultSetColumnType(
+ name="field",
+ column_name="field",
+ type=datatype.parse_sqltype(type_str),
+ is_dttm=False,
+ )
+
+ expanded = TrinoEngineSpec._expand_columns(col)
+
+ assert len(expanded) == depth + 1
+ assert expanded[-1]["name"] == "field" + ".a" * depth
+
+
+def test_expand_columns_has_no_recursion_depth_guard() -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when a depth guard lands.
+
+ ``_expand_columns`` recurses once per ROW level with no maximum depth, so a
+ type nested deeper than the interpreter's recursion limit raises
+ ``RecursionError`` and would break column inspection for the whole table.
+ The type is built programmatically because ``parse_sqltype`` is itself
+ recursive and would fail first on a type string this deep.
+ """
+ import sys
+
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ col_type: Any = types.String()
+ for _ in range(sys.getrecursionlimit() + 100):
+ col_type = datatype.ROW([("a", col_type)])
+ col = ResultSetColumnType(
+ name="field", column_name="field", type=col_type, is_dttm=False
+ )
+
+ with pytest.raises(RecursionError):
+ TrinoEngineSpec._expand_columns(col)
+
+
+# ---------------------------------------------------------------------------
+# Error mapping.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("exception", [TrinoConnectionError, TrinoAuthError])
+def test_get_dbapi_exception_mapping_matches_subclasses(
+ exception: type[Exception],
+) -> None:
+ """
+ The mapping matches by ``issubclass``, so driver exceptions without a
+ bespoke entry (connection and auth failures subclass ``OperationalError``)
+ still map to ``SupersetDBAPIOperationalError``.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) == SupersetDBAPIOperationalError
+
+
+@pytest.mark.parametrize("exception", [Http502Error, IntegrityError])
+def test_get_dbapi_exception_mapping_ignores_unmapped_exceptions(
+ exception: type[Exception],
+) -> None:
+ """
+ Exceptions outside the three mapped categories fall through to the caller's
+ default: Trino's HTTP 5xx errors subclass ``HttpError`` (not a
+ ``DatabaseError``), and ``IntegrityError`` subclasses ``DatabaseError``
+ directly rather than one of the handled subclasses.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapping = TrinoEngineSpec.get_dbapi_exception_mapping()
+
+ assert mapping.get(exception) is None
+ assert mapping.get(exception, SupersetDBAPIConnectionError) is (
+ SupersetDBAPIConnectionError
+ )
+
+
+@pytest.mark.parametrize(
+ "exception,expected_type",
+ [
+ (TrinoUserError({"message": "syntax error"}), SupersetDBAPIProgrammingError),
+ (TrinoInternalError({"message": "oops"}), SupersetDBAPIDatabaseError),
+ (TrinoExternalError({"message": "hive down"}), SupersetDBAPIOperationalError),
+ ],
+)
+def test_get_dbapi_mapped_exception_keeps_message(
+ exception: Exception, expected_type: type[Exception]
+) -> None:
+ """
+ The mapped Superset exception is built from the driver exception's text,
+ so the original reason stays visible to the user.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mapped = TrinoEngineSpec.get_dbapi_mapped_exception(exception)
+
+ assert type(mapped) is expected_type
+ assert str(mapped) == str(exception)
+
+
+def test_trino_has_no_custom_error_patterns() -> None:
+ """
+ Characterizes current behavior (production gap). Invert if Trino gains
+ ``custom_errors``.
+
+ Unlike ``PrestoEngineSpec``, which rewrites nine known Presto errors into
+ friendly typed messages, Trino defines no ``custom_errors``: every driver
+ error, including a missing table or bad credentials, surfaces as a generic
+ database engine error carrying the raw Trino text.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ assert TrinoEngineSpec.custom_errors == {}
+
+ errors = TrinoEngineSpec.extract_errors(
+ TrinoUserError({"message": "line 1:15: Table 'tpch.sf1.nope' does not exist"})
+ )
+
+ assert len(errors) == 1
+ assert errors[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert "Table 'tpch.sf1.nope' does not exist" in errors[0].message
+ assert errors[0].extra is not None
+ assert errors[0].extra["engine_name"] == "Trino"
+
+
+# ---------------------------------------------------------------------------
+# Function names and extra table metadata.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "functions,expected",
+ [
+ (["abs", "array_agg", "json_extract"], ["abs", "array_agg", "json_extract"]),
+ ([], []),
+ ],
+)
+def test_get_function_names(
+ mocker: MockerFixture, functions: list[str], expected: list[str]
+) -> None:
+ """
+ get_function_names is defined on PrestoBaseEngineSpec; mirrors
+ test_presto.py::test_get_function_names_lists_presto_functions at the Trino
+ layer.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": functions})
+
+ assert TrinoEngineSpec.get_function_names(database) == expected
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_database_function_names_swallows_connection_error(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ ``Database.function_names`` feeds SQL Lab autocomplete and must not crash
+ when Trino is unreachable: the failure is logged and an empty list returned
+ (https://github.com/apache/superset/issues/9678).
+ """
+ database, _ = trino_database
+ mocker.patch.object(
+ database, "get_df", side_effect=TrinoConnectionError("connection refused")
+ )
+
+ assert database.function_names == []
+
+
+def test_get_extra_table_metadata_no_partitions(mocker: MockerFixture) -> None:
+ """
+ A table with no indexes at all (not just Iceberg metadata filtered out)
+ yields no partition metadata and no latest-partition query.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ db_mock = mocker.MagicMock()
+ db_mock.get_indexes = Mock(return_value=[])
+ db_mock.get_extra = Mock(return_value={})
+ db_mock.has_view = Mock(return_value=False)
+ db_mock.get_df = Mock()
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ db_mock, Table("test_table", "test_schema")
+ )
+
+ assert result == {}
+ db_mock.get_df.assert_not_called()
+
+
+def test_get_extra_table_metadata_for_dropped_table(
+ mocker: MockerFixture,
+ trino_database: tuple[Any, MagicMock],
+) -> None:
+ """
+ A table dropped between listing and inspection raises ``NoSuchTableError``
+ from the inspector. ``TrinoEngineSpec.get_indexes`` turns that into an empty
+ index list, so ``get_extra_table_metadata`` returns no metadata instead of
+ failing; exercised through a real ``Database`` so the production handler
+ (not a mock of it) is what runs.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ database, inspector = trino_database
+ inspector.get_indexes.side_effect = NoSuchTableError(
+ "The specified table does not exist."
+ )
+ mocker.patch.object(database, "has_view", return_value=False)
+
+ result = TrinoEngineSpec.get_extra_table_metadata(
+ database, Table("dropped_table", "test_schema")
+ )
+
+ assert result == {}
+
+
+# ---------------------------------------------------------------------------
+# Query cancellation, threading, and OAuth2.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "malformed_id",
+ [
+ "123; DROP TABLE users",
+ "123' OR '1'='1",
+ "../../etc/passwd",
+ "123 UNION SELECT 1",
+ "abc-def",
+ "",
+ ],
+)
+def test_cancel_query_rejects_malformed_query_id(malformed_id: str) -> None:
+ """
+ A malformed or injection-style ``cancel_query_id`` is rejected before it can
+ be interpolated into the ``CALL system.runtime.kill_query`` statement:
+ ``cancel_query`` returns False and never touches the cursor.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.models.sql_lab import Query
+
+ cursor_mock = Mock()
+
+ assert TrinoEngineSpec.cancel_query(cursor_mock, Query(), malformed_id) is False
+ cursor_mock.execute.assert_not_called()
+
+
+def test_execute_with_cursor_shares_mutable_g_state_across_threads(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert
+ when the copy is made isolated.
+
+ The worker thread rebuilds Flask's ``g`` with ``setattr(g, key, value)``
+ per attribute: a shallow copy. A mutable attribute (here a dict) is the
+ same object on both sides, so a mutation made in the worker thread is
+ visible to the request thread.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "existing-query-id"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ def _mutate_in_worker(*args: Any, **kwargs: Any) -> None:
+ g.shared_state["count"] += 1
+
+ with app.test_request_context("/some/place/"):
+ shared_state = {"count": 0}
+ g.shared_state = shared_state
+
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=_mutate_in_worker),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert shared_state["count"] == 1
+
+
+@pytest.mark.parametrize("exception", [RuntimeError("boom"), MemoryError("oom")])
+def test_execute_with_cursor_reraises_worker_exception_before_query_id(
+ app, mocker: MockerFixture, exception: Exception
+) -> None:
+ """
+ A non-DB exception raised by the worker before Trino assigns a query ID is
+ re-raised unchanged in the calling thread. The wait for a query ID ends
+ because the worker's ``finally`` sets the execute event, so ``handle_cursor``
+ is still called (with no query ID yet) instead of spinning forever.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=exception):
+ with pytest.raises(type(exception)) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is exception
+ assert observed_query_ids == [None]
+
+
+def test_execute_with_cursor_reraises_worker_exception_after_query_id(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ When the query ID became available before the worker failed, progress
+ handling sees that real ID and the worker's exception is still re-raised
+ in the calling thread once the cursor has been handled.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = None
+ observed_query_ids: list[str | None] = []
+ mocker.patch.object(
+ TrinoEngineSpec,
+ "handle_cursor",
+ side_effect=lambda cursor, query: observed_query_ids.append(cursor.query_id),
+ )
+ failure = RuntimeError("boom")
+
+ def _fail_after_query_id(*args: Any, **kwargs: Any) -> None:
+ mock_cursor.query_id = "myQueryId"
+ raise failure
+
+ with app.test_request_context("/some/place/"):
+ with patch.object(TrinoEngineSpec, "execute", side_effect=_fail_after_query_id):
+ with pytest.raises(RuntimeError) as exc_info:
+ TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert exc_info.value is failure
+ assert observed_query_ids == ["myQueryId"]
+
+
+@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
+def test_execute_with_cursor_swallows_base_exception_from_worker(
+ app, mocker: MockerFixture
+) -> None:
+ """
+ Characterizes current behavior (production defect). Invert when fixed.
+
+ The worker only records ``Exception`` subclasses. A ``BaseException`` such
+ as ``KeyboardInterrupt`` or ``SystemExit`` escapes the ``except`` clause, so
+ nothing is stored in the result dict and ``execute_with_cursor`` returns
+ normally as though the query had succeeded.
+ """
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+
+ mock_cursor = mocker.MagicMock()
+ mock_cursor.query_id = "myQueryId"
+ mocker.patch.object(TrinoEngineSpec, "handle_cursor")
+
+ with app.test_request_context("/some/place/"):
+ with (
+ patch.object(TrinoEngineSpec, "execute", side_effect=KeyboardInterrupt),
+ patch.dict("superset.config.DISALLOWED_SQL_FUNCTIONS", {}, clear=True),
+ ):
+ result = TrinoEngineSpec.execute_with_cursor(
+ cursor=mock_cursor,
+ sql="SELECT 1 FROM foo",
+ query=mocker.MagicMock(),
+ )
+
+ assert result is None
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and marked this thread resolved** as of `86cc479`.
The worker now mutates `g.shared_state`, and the test is marked `xfail(strict=True, raises=AssertionError)` while asserting the caller’s state remains unchanged. It no longer treats a worker `KeyboardInterrupt` as successful execution.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it open.</sub>
<!-- codeant-auto-resolve-reply -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org