This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 669579e6c84 Fix SnowflakeSqlApiOperator polling to avoid needless
sleeps (#69450)
669579e6c84 is described below
commit 669579e6c844dc6b73d2acacb411b1e57b860eb9
Author: Amogh Desai <[email protected]>
AuthorDate: Mon Jul 6 18:34:26 2026 +0530
Fix SnowflakeSqlApiOperator polling to avoid needless sleeps (#69450)
---
.../providers/snowflake/operators/snowflake.py | 8 +++-
.../unit/snowflake/operators/test_snowflake.py | 51 ++++++++++++++++++++--
2 files changed, 53 insertions(+), 6 deletions(-)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
index eec2e39ecc2..b9e74b18589 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 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
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 cea968c38cd..b1d6b418839 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 @@ class TestSnowflakeSqlApiOperator:
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 @@ class TestSnowflakeSqlApiOperator:
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
