diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py index 68d1856e18c..e74d78a2268 100644 --- a/superset/db_engine_specs/trino.py +++ b/superset/db_engine_specs/trino.py @@ -225,7 +225,9 @@ class TrinoEngineSpec(PrestoBaseEngineSpec): terminal_states = {"FINISHED", "FAILED", "CANCELED"} state = "QUEUED" progress = 0.0 - poll_interval = current_app.config["DB_POLL_INTERVAL_SECONDS"].get(cls.engine, 1) + poll_interval = current_app.config["DB_POLL_INTERVAL_SECONDS"].get( + cls.engine, 1 + ) max_wait_time = current_app.config.get("SQLLAB_ASYNC_TIME_LIMIT_SEC", 21600) start_time = time.time() while state not in terminal_states: diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py index 04703dd567e..b8418201e03 100644 --- a/tests/unit_tests/db_engine_specs/test_trino.py +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -386,7 +386,7 @@ def test_prepare_cancel_query( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_early_cancel( mock_app: Mock, mock_db: Mock, @@ -924,7 +924,7 @@ def test_timegrain_expressions(time_grain: str, expected_result: str) -> None: @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_progress_updates( mock_app: Mock, mock_db: Mock, @@ -974,7 +974,7 @@ def test_handle_cursor_progress_updates( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_cancels_on_stopped_status( mock_app: Mock, mock_db: Mock, @@ -1009,7 +1009,7 @@ def test_handle_cursor_cancels_on_stopped_status( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_cancels_on_timed_out_status( mock_app: Mock, mock_db: Mock, @@ -1044,7 +1044,7 @@ def test_handle_cursor_cancels_on_timed_out_status( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_breaks_on_execute_error( mock_app: Mock, mock_db: Mock, @@ -1079,7 +1079,7 @@ def test_handle_cursor_breaks_on_execute_error( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_breaks_on_execute_event_set( mock_app: Mock, mock_db: Mock, @@ -1120,7 +1120,7 @@ def test_handle_cursor_breaks_on_execute_event_set( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_handles_zero_total_splits( mock_app: Mock, mock_db: Mock, @@ -1165,7 +1165,7 @@ def test_handle_cursor_handles_zero_total_splits( @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") +@patch("superset.db_engine_specs.trino.current_app") def test_handle_cursor_only_commits_on_progress_change( mock_app: Mock, mock_db: Mock,