drivaspreset commented on code in PR #43607:
URL: https://github.com/apache/superset/pull/43607#discussion_r3883110488
##########
tests/unit_tests/db_engine_specs/test_presto.py:
##########
@@ -412,6 +656,635 @@ def test_extract_errors_maps_401_to_access_denied() ->
None:
assert result[0].error_type ==
SupersetErrorType.CONNECTION_ACCESS_DENIED_ERROR
[email protected](
+ "raw_message,context,expected_error_type,expected_message",
+ [
+ pytest.param(
+ "line 1:8: Column 'bar' cannot be resolved",
+ {},
+ "COLUMN_DOES_NOT_EXIST_ERROR",
+ 'We can\'t seem to resolve the column "bar" at line 1:8.',
+ id="column_does_not_exist",
+ ),
+ pytest.param(
+ "Table 'default.foo' does not exist",
+ {},
+ "TABLE_DOES_NOT_EXIST_ERROR",
+ "The table \"'default.foo'\" does not exist. "
+ "A valid table must be used to run this query.",
+ id="table_does_not_exist",
+ ),
+ pytest.param(
+ "line 1:15: Schema 'bar' does not exist",
+ {},
+ "SCHEMA_DOES_NOT_EXIST_ERROR",
+ 'The schema "bar" does not exist. '
+ "A valid schema must be used to run this query.",
+ id="schema_does_not_exist",
+ ),
+ pytest.param(
+ "Access Denied: Invalid credentials",
+ {"username": "bob"},
+ "CONNECTION_ACCESS_DENIED_ERROR",
+ 'Either the username "bob" or the password is incorrect.',
+ id="access_denied_invalid_credentials",
+ ),
+ pytest.param(
+ "presto error: Unexpected status code 401 b'Unauthorized'",
+ {},
+ "CONNECTION_ACCESS_DENIED_ERROR",
+ "Unexpected HTTP 401 response. Check your credentials.",
+ id="access_denied_http_401",
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 8] nodename nor "
+ "servname provided, or not known",
+ {"hostname": "badhost"},
+ "CONNECTION_INVALID_HOSTNAME_ERROR",
+ 'The hostname "badhost" cannot be resolved.',
+ id="invalid_hostname",
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 60] Operation timed
out",
+ {"hostname": "myhost", "port": 8080},
+ "CONNECTION_HOST_DOWN_ERROR",
+ 'The host "myhost" might be down, and can\'t be reached on port
8080.',
+ id="host_down_operation_timed_out",
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 61] Connection
refused",
+ {"hostname": "myhost", "port": 8080},
+ "CONNECTION_PORT_CLOSED_ERROR",
+ 'Port 8080 on hostname "myhost" refused the connection.',
+ id="port_closed",
+ ),
+ pytest.param(
+ "line 1:8: Catalog 'foo' does not exist",
+ {},
+ "CONNECTION_UNKNOWN_DATABASE_ERROR",
+ 'Unable to connect to catalog named "foo".',
+ id="unknown_catalog",
+ ),
+ ],
+)
+def test_extract_errors_matches_all_custom_error_patterns(
+ raw_message: str,
+ context: dict[str, Any],
+ expected_error_type: str,
+ expected_message: str,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+ from superset.errors import ErrorLevel, SupersetErrorType
+
+ result = PrestoEngineSpec.extract_errors(Exception(raw_message),
context=context)
+
+ assert len(result) == 1
+ assert result[0].error_type == getattr(SupersetErrorType,
expected_error_type)
+ assert result[0].message == expected_message
+ assert result[0].level == ErrorLevel.ERROR
+ assert result[0].extra is not None
+ assert result[0].extra["engine_name"] == "Presto"
+
+
[email protected](
+ "raw_message,missing_placeholder",
+ [
+ pytest.param(
+ "Access Denied: Invalid credentials", "username",
id="access_denied"
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 8] nodename nor "
+ "servname provided, or not known",
+ "hostname",
+ id="invalid_hostname",
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 60] Operation timed
out",
+ "hostname",
+ id="host_down",
+ ),
+ pytest.param(
+ "Failed to establish a new connection: [Errno 61] Connection
refused",
+ "port",
+ id="port_closed",
+ ),
+ ],
+)
+def test_extract_errors_raises_key_error_without_context(
+ raw_message: str,
+ missing_placeholder: str,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ with pytest.raises(KeyError, match=missing_placeholder):
+ PrestoEngineSpec.extract_errors(Exception(raw_message))
+
+
+def test_extract_errors_returns_first_matching_pattern() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+ from superset.errors import SupersetErrorType
+
+ msg = "line 1:8: Table 'x' does not exist and Column 'bar' cannot be
resolved"
+ result = PrestoEngineSpec.extract_errors(Exception(msg))
+
+ assert len(result) == 1
+ assert result[0].error_type ==
SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR
+
+
+def test_extract_errors_falls_back_to_generic_error() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+ from superset.errors import ErrorLevel, SupersetErrorType
+
+ result = PrestoEngineSpec.extract_errors(Exception("Generic Error"))
+
+ assert len(result) == 1
+ assert result[0].error_type == SupersetErrorType.GENERIC_DB_ENGINE_ERROR
+ assert result[0].message == "Generic Error"
+ assert result[0].level == ErrorLevel.ERROR
+ assert result[0].extra is not None
+ assert result[0].extra["engine_name"] == "Presto"
+ assert result[0].extra["issue_codes"] == [
+ {
+ "code": 1002,
+ "message": "Issue 1002 - The database returned an unexpected
error.",
+ }
+ ]
+
+
+def test_extract_error_message_from_orig_database_error() -> None:
+ from collections import namedtuple
+
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ DatabaseError = namedtuple("DatabaseError", ["error_dict"]) # noqa: N806
+ db_err = DatabaseError(
+ {"errorName": "name", "errorLocation": "location", "message": "msg"}
+ )
+ exception = Exception()
+ exception.orig = db_err # type: ignore[attr-defined]
+
+ assert PrestoEngineSpec._extract_error_message(exception) == "name at
location: msg"
+
+
+def test_extract_error_message_from_database_error_args() -> None:
+ from pyhive.exc import DatabaseError
+
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ exception = DatabaseError({"message": "Err message"})
+
+ assert PrestoEngineSpec._extract_error_message(exception) == "Err message"
+
+
+def test_extract_error_message_from_database_error_without_message() -> None:
+ from pyhive.exc import DatabaseError
+
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ exception = DatabaseError({"errorName": "SYNTAX_ERROR"})
+
+ assert str(PrestoEngineSpec._extract_error_message(exception)) == (
+ "Unknown Presto Error"
+ )
+
+
+def test_extract_error_message_from_general_exception() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ assert (
+ PrestoEngineSpec._extract_error_message(Exception("Err message"))
+ == "Err message"
+ )
+
+
+def test_expand_data_returns_input_untouched_when_flag_disabled() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ columns: list[ResultSetColumnType] = [
+ {
+ "column_name": "row_column",
+ "name": "row_column",
+ "type": "ROW(NESTED_OBJ VARCHAR)",
+ "is_dttm": False,
+ }
+ ]
+ data = [{"row_column": ["a"]}]
+
+ result_columns, result_data, expanded =
PrestoEngineSpec.expand_data(columns, data)
+
+ assert result_columns is columns
+ assert result_data is data
+ assert expanded == []
+
+
[email protected](
+ "superset.extensions.feature_flag_manager._feature_flags",
+ {"PRESTO_EXPAND_DATA": True},
+ clear=True,
+)
+def test_expand_data_flattens_deeply_nested_row_columns() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ columns: list[ResultSetColumnType] = [
+ {
+ "column_name": "r",
+ "name": "r",
+ "type": "ROW(L1 ROW(L2 ROW(L3 ROW(L4 VARCHAR))))",
+ "is_dttm": False,
+ }
+ ]
+ data = [{"r": [[[["deep"]]]]}]
+
+ result_columns, result_data, expanded =
PrestoEngineSpec.expand_data(columns, data)
+
+ assert [column["column_name"] for column in result_columns] == [
+ "r",
+ "r.l1",
+ "r.l1.l2",
+ "r.l1.l2.l3",
+ "r.l1.l2.l3.l4",
+ ]
+ assert [column["column_name"] for column in expanded] == [
+ "r.l1",
+ "r.l1.l2",
+ "r.l1.l2.l3",
+ "r.l1.l2.l3.l4",
+ ]
+ assert result_data[0]["r.l1.l2.l3.l4"] == "deep"
+
+
[email protected](
+ "superset.extensions.feature_flag_manager._feature_flags",
+ {"PRESTO_EXPAND_DATA": True},
+ clear=True,
+)
+def test_expand_data_raises_on_malformed_json_in_array_column() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ columns: list[ResultSetColumnType] = [
+ {
+ "column_name": "array_column",
+ "name": "array_column",
+ "type": "ARRAY(BIGINT)",
+ "is_dttm": False,
+ }
+ ]
+
+ with pytest.raises(ValueError, match="Expecting value"):
+ PrestoEngineSpec.expand_data(columns, [{"array_column": "not json"}])
+
+
[email protected](
+ "superset.extensions.feature_flag_manager._feature_flags",
+ {"PRESTO_EXPAND_DATA": True},
+ clear=True,
+)
+def test_expand_data_raises_on_malformed_json_in_row_column() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ columns: list[ResultSetColumnType] = [
+ {
+ "column_name": "row_column",
+ "name": "row_column",
+ "type": "ROW(NESTED_OBJ VARCHAR)",
+ "is_dttm": False,
+ }
+ ]
+
+ with pytest.raises(ValueError, match="Expecting value"):
+ PrestoEngineSpec.expand_data(columns, [{"row_column": "not json"}])
+
+
+def test_get_function_names_lists_presto_functions() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mock.MagicMock()
+ database.get_df.return_value = pd.DataFrame(
+ {"Function": ["abs", "avg", "cardinality"]}
+ )
+
+ assert PrestoEngineSpec.get_function_names(database) == [
+ "abs",
+ "avg",
+ "cardinality",
+ ]
+ database.get_df.assert_called_once_with("SHOW FUNCTIONS")
+
+
+def test_get_function_names_returns_empty_list_for_no_functions() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mock.MagicMock()
+ database.get_df.return_value = pd.DataFrame({"Function": []})
+
+ assert PrestoEngineSpec.get_function_names(database) == []
+
+
+def test_get_function_names_raises_on_dataframe_without_function_column() ->
None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mock.MagicMock()
+ database.get_df.return_value = pd.DataFrame()
+
+ with pytest.raises(KeyError, match="Function"):
+ PrestoEngineSpec.get_function_names(database)
+
+
+def test_get_function_names_propagates_connection_error() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mock.MagicMock()
+ database.get_df.side_effect = Exception("Connection refused")
+
+ with pytest.raises(Exception, match="Connection refused"):
+ PrestoEngineSpec.get_function_names(database)
+
+
[email protected](
+ "extra,expected",
+ [
+ pytest.param({}, False, id="no_version_key"),
+ pytest.param({"version": None}, False, id="version_none"),
+ pytest.param({"version": "0.318"}, False, id="just_below_gate"),
+ pytest.param({"version": "0.319"}, True, id="exactly_at_gate"),
+ pytest.param({"version": "0.400"}, True, id="above_gate"),
+ ],
+)
+def test_get_allow_cost_estimate_version_gate(
+ extra: dict[str, Any],
+ expected: bool,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ assert PrestoEngineSpec.get_allow_cost_estimate(extra) is expected
+
+
+def test_get_allow_cost_estimate_rejects_unparseable_version() -> None:
+ from packaging.version import InvalidVersion
+
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ with pytest.raises(InvalidVersion):
+ PrestoEngineSpec.get_allow_cost_estimate({"version": "not-a-version"})
+
+
+def test_estimate_statement_cost() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ cursor = mock.MagicMock()
+ cursor.fetchone.return_value = ['{"a": "b"}']
+
+ result = PrestoEngineSpec.estimate_statement_cost(
+ mock.MagicMock(), "SELECT * FROM birth_names", cursor
+ )
+
+ assert result == {"a": "b"}
+ cursor.execute.assert_called_once_with(
+ "EXPLAIN (TYPE IO, FORMAT JSON) SELECT * FROM birth_names"
+ )
+
+
+def test_estimate_statement_cost_propagates_execute_failure() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ cursor = mock.MagicMock()
+ cursor.execute.side_effect = Exception("line 1:1: mismatched input 'DROP'")
+
+ with pytest.raises(Exception, match="mismatched input"):
+ PrestoEngineSpec.estimate_statement_cost(
+ mock.MagicMock(), "DROP TABLE birth_names", cursor
+ )
+
+
+def test_query_cost_formatter() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ raw_cost = [
+ {
+ "estimate": {
+ "outputRowCount": 9.04969899e8,
+ "outputSizeInBytes": 3.54143678301e11,
+ "cpuCost": 3.54143678301e11,
+ "maxMemory": 0.0,
+ "networkCost": 3.54143678301e11,
+ },
+ }
+ ]
+
+ assert PrestoEngineSpec.query_cost_formatter(raw_cost) == [
+ {
+ "Output count": "904 M rows",
+ "Output size": "354 GB",
+ "CPU cost": "354 G",
+ "Max memory": "0 B",
+ "Network cost": "354 G",
+ }
+ ]
+
+
+def test_query_cost_formatter_omits_missing_estimate_keys() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ raw_cost = [{"estimate": {"outputRowCount": 1234.0}}, {}]
+
+ assert PrestoEngineSpec.query_cost_formatter(raw_cost) == [
+ {"Output count": "1 K rows"},
+ {},
+ ]
+
+
+def test_query_cost_formatter_raises_on_null_estimate_value() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ with pytest.raises(TypeError):
+ PrestoEngineSpec.query_cost_formatter(
+ [{"estimate": {"outputRowCount": None, "outputSizeInBytes": 1.0}}]
+ )
+
+
+def test_estimate_query_cost_raises_when_version_too_old(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_extra.return_value = {"version": "0.318"}
+
+ with pytest.raises(Exception, match="Database does not support cost
estimation"):
+ PrestoEngineSpec.estimate_query_cost(
+ database, "hive", "default", "SELECT 1", None
+ )
+
+ database.get_raw_connection.assert_not_called()
+
+
+def test_estimate_query_cost_estimates_each_statement(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ database = mocker.MagicMock()
+ database.get_extra.return_value = {"version": "0.400"}
+ database.mutate_sql_based_on_config.side_effect = lambda sql, **_kwargs:
sql
+ cursor = mock.MagicMock()
+ cursor.fetchone.side_effect = [
+ ['{"estimate": {"outputRowCount": 1.0}}'],
+ ['{"estimate": {"outputRowCount": 2.0}}'],
+ ]
+
database.get_raw_connection.return_value.__enter__.return_value.cursor.return_value
= ( # noqa: E501
+ cursor
+ )
+
+ result = PrestoEngineSpec.estimate_query_cost(
+ database, "hive", "default", "SELECT 1; SELECT 2", None
+ )
+
+ assert result == [
+ {"estimate": {"outputRowCount": 1.0}},
+ {"estimate": {"outputRowCount": 2.0}},
+ ]
+ assert cursor.execute.call_args_list == [
+ mock.call("EXPLAIN (TYPE IO, FORMAT JSON) SELECT\n 1"),
+ mock.call("EXPLAIN (TYPE IO, FORMAT JSON) SELECT\n 2"),
+ ]
+
+
+TRACKING_URL = (
+ "https://presto.example.com:8080/ui/query.html?20220101_120000_00001_abcde";
+)
+
+
+def _presto_cursor() -> mock.MagicMock:
+ cursor = mock.MagicMock()
+ cursor._protocol = "https"
+ cursor._host = "presto.example.com"
+ cursor._port = 8080
+ cursor.last_query_id = "20220101_120000_00001_abcde"
+ return cursor
+
+
+def _handle_cursor_query(
+ mocker: MockerFixture,
+) -> tuple[mock.MagicMock, mock.MagicMock]:
+ from superset.common.db_query_status import QueryStatus
+
+ mock_db = mocker.patch("superset.db_engine_specs.presto.db")
+ query = mock.MagicMock()
+ query.id = 42
+ query.progress = 0
+ query.status = QueryStatus.RUNNING
+ query.database.connect_args = {"poll_interval": 0}
+ mock_db.session.query.return_value.filter_by.return_value.one.return_value
= query
+ return mock_db, query
+
+
+def test_get_tracking_url_builds_presto_ui_link() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ assert PrestoEngineSpec.get_tracking_url(_presto_cursor()) == TRACKING_URL
+
+
+def test_get_tracking_url_returns_none_for_falsy_query_id() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ cursor = _presto_cursor()
+ cursor.last_query_id = None
+
+ assert PrestoEngineSpec.get_tracking_url(cursor) is None
+
+
+def test_get_tracking_url_returns_none_when_attribute_absent() -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ assert PrestoEngineSpec.get_tracking_url(mock.Mock(spec=[])) is None
+
+
+def test_handle_cursor_records_tracking_url_and_progress(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ mock_db, query = _handle_cursor_query(mocker)
+ cursor = _presto_cursor()
+ cursor.poll.side_effect = [
+ {"stats": {"state": "RUNNING", "completedSplits": 5, "totalSplits":
10}},
+ None,
+ ]
+
+ PrestoEngineSpec.handle_cursor(cursor, query)
+
+ assert query.tracking_url == TRACKING_URL
+ assert query.progress == 50.0
+ assert cursor.poll.call_count == 2
+ cursor.cancel.assert_not_called()
+ assert mock_db.session.commit.called
+
+
+def test_handle_cursor_stops_polling_when_query_finished(
+ mocker: MockerFixture,
+) -> None:
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ _mock_db, query = _handle_cursor_query(mocker)
+ cursor = _presto_cursor()
+ cursor.poll.side_effect = [{"stats": {"state": "FINISHED"}}]
+
+ PrestoEngineSpec.handle_cursor(cursor, query)
+
+ assert query.progress == 0
+ assert cursor.poll.call_count == 1
+ cursor.cancel.assert_not_called()
+
+
[email protected]("status", ["STOPPED", "TIMED_OUT"])
+def test_handle_cursor_cancels_when_user_stops_query(
+ mocker: MockerFixture,
+ status: str,
+) -> None:
+ from superset.common.db_query_status import QueryStatus
+ from superset.db_engine_specs.presto import PrestoEngineSpec
+
+ _mock_db, query = _handle_cursor_query(mocker)
+ query.status = getattr(QueryStatus, status)
Review Comment:
Yes you're right about the boundary. I confirmed it by mutation: with the
db.session.query(...).filter_by(id=query_id).one() reload removed from
handle_cursor, the old test still passed
Updated the test: the query passed in stays RUNNING, the mocked session
returns a distinct STOPPED/TIMED_OUT object, and the test asserts the reload
happened with the right id and that it preceded cursor.cancel(). With the
reload removed it now fails on both statuses
--
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]
