Fix SnowflakeSqlApiOperator polling to avoid needless sleeps

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,

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