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).
+# ---------------------------------------------------------------------------
+
+
[email protected]
+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)
+
+
[email protected]("method", ["get_table_names", "get_view_names"])
[email protected](
+ "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)
+
+
[email protected]("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")
+
+
[email protected]("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)
+
+
[email protected](
+ "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
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected]("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
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected]("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
+
+
[email protected]("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
+ )
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "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
+
+
[email protected]("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"]
+
+
[email protected]("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).
+# ---------------------------------------------------------------------------
+
+
[email protected]
+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)
+
+
[email protected]("method", ["get_table_names", "get_view_names"])
[email protected](
+ "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)
+
+
[email protected]("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")
+
+
[email protected]("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)
+
+
[email protected](
+ "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
+
+
[email protected](
+ "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)
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected]("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
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected]("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
+
+
[email protected]("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
+ )
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "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.
+# ---------------------------------------------------------------------------
+
+
[email protected](
+ "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
+
+
[email protected]("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"]
+
+
[email protected]("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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
