diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index eec2e39ecc2c6..b9e74b185890a 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -491,11 +491,15 @@ def poll_on_queries(self): if statement_status.get("status") == "error": queries_in_progress.remove(query_id) statement_error_status[query_id] = statement_status - if statement_status.get("status") == "success": + elif statement_status.get("status") == "success": statement_success_status[query_id] = statement_status queries_in_progress.remove(query_id) - if statement_status.get("status") == "running": + elif statement_status.get("status") == "running": statement_running_status[query_id] = statement_status + # Only wait before the next poll cycle if something is still running. Sleeping + # unconditionally after every handle would delay returning even when this cycle + # already resolved everything (e.g. all statements finished, or one failed). + if queries_in_progress: time.sleep(self.poll_interval) return { "success": statement_success_status, diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py index cea968c38cd88..b1d6b4188393a 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -359,6 +359,47 @@ def test_poll_on_queries_raises_runtime_error_on_status_check_failure( with pytest.raises(RuntimeError, match="Failed to get status for query uuid1"): operator.poll_on_queries() + def test_poll_on_queries_no_sleep_when_all_resolved(self, mock_get_sql_api_query_status): + operator = SnowflakeSqlApiOperator( + task_id=TASK_ID, + snowflake_conn_id="snowflake_default", + sql=SQL_MULTIPLE_STMTS, + statement_count=4, + do_xcom_push=False, + ) + operator.query_ids = ["uuid1", "uuid2"] + mock_get_sql_api_query_status.side_effect = [{"status": "success"}, {"status": "error"}] + + with mock.patch("time.sleep") as mock_sleep: + result = operator.poll_on_queries() + + mock_sleep.assert_not_called() + assert result["success"] == {"uuid1": {"status": "success"}} + assert result["error"] == {"uuid2": {"status": "error"}} + assert result["running"] == {} + + def test_poll_on_queries_sleeps_once_per_cycle(self, mock_get_sql_api_query_status): + """One handle is still running, so the cycle sleeps -- but only once, not per handle.""" + operator = SnowflakeSqlApiOperator( + task_id=TASK_ID, + snowflake_conn_id="snowflake_default", + sql=SQL_MULTIPLE_STMTS, + statement_count=4, + do_xcom_push=False, + ) + operator.query_ids = ["uuid1", "uuid2", "uuid3"] + mock_get_sql_api_query_status.side_effect = [ + {"status": "success"}, + {"status": "running"}, + {"status": "success"}, + ] + + with mock.patch("time.sleep") as mock_sleep: + result = operator.poll_on_queries() + + mock_sleep.assert_called_once_with(operator.poll_interval) + assert result["running"] == {"uuid2": {"status": "running"}} + @pytest.mark.parametrize( ("mock_sql", "statement_count"), [pytest.param(SQL_MULTIPLE_STMTS, 4, id="multi"), pytest.param(SINGLE_STMT, 1, id="single")], @@ -581,18 +622,20 @@ def test_snowflake_sql_api_execute_operator_polling_running( mock_get_sql_api_query_status.side_effect = [ # Initial get_sql_api_query_status check {"status": "running"}, - # 1st poll_on_queries check (poll_interval: 5s) + # 1st poll_on_queries check (poll_interval: 5s) -- still running, sleeps {"status": "running"}, - # 2nd poll_on_queries check (poll_interval: 5s) + # 2nd poll_on_queries check (poll_interval: 5s) -- still running, sleeps {"status": "running"}, - # 3rd poll_on_queries check (poll_interval: 5s) + # 3rd poll_on_queries check -- resolves to success, no sleep needed {"status": "success"}, ] with mock.patch("time.sleep") as mock_sleep: operator.execute(context=None) mock_check_query_output.assert_called_once_with(["uuid1"]) - assert mock_sleep.call_count == 3 + # Only 2 sleeps: the cycle that resolves the last running query returns + # immediately instead of sleeping once more before reporting success. + assert mock_sleep.call_count == 2 def test_snowflake_sql_api_execute_operator_polling_failed( self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output