This is an automated email from the ASF dual-hosted git repository.
potiuk 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 3bd4df6b2da Fix Snowflake SQL API wait_for_query timing out on
finished queries (#73761)
3bd4df6b2da is described below
commit 3bd4df6b2da90320412c486ccc2ad861b33b9b8a
Author: Karthik Mohankumar <[email protected]>
AuthorDate: Sun Oct 4 19:22:12 2026 -0500
Fix Snowflake SQL API wait_for_query timing out on finished queries (#73761)
* Fix Snowflake SQL API wait_for_query timing out on finished queries
wait_for_query checked the elapsed time before looking at the status it had
just fetched, so a status call that returned after the timeout raised
TimeoutError even when the query had already succeeded or failed. Status calls
can take several seconds because they retry with exponential backoff, and the
OpenLineage integration waits with a 3-second timeout, so a finished query
could be reported as timed out and its lineage details dropped. Elapsed time
now also uses the monotonic clock [...]
* Test wait_for_query with a status call that outlasts the timeout
---
.../providers/snowflake/hooks/snowflake_sql_api.py | 12 ++++----
.../unit/snowflake/hooks/test_snowflake_sql_api.py | 35 +++++++++++++++++++---
2 files changed, 37 insertions(+), 10 deletions(-)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
index 05d80a39e48..4dca2b0a12a 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
@@ -383,21 +383,21 @@ class SnowflakeSqlApiHook(SnowflakeHook):
:raises RuntimeError: If the query status is 'error' and `raise_error`
is True.
:raises TimeoutError: If the query doesn't finish within the specified
timeout.
"""
- start_time = time.time()
+ start_time = time.monotonic()
while True:
response = self.get_sql_api_query_status(query_id=query_id)
self.log.debug("Query status `%s`", response["status"])
- if time.time() - start_time > timeout:
- raise TimeoutError(
- f"Query `{query_id}` did not finish within the timeout
period of {timeout} seconds."
- )
-
if response["status"] != "running":
self.log.info("Query status `%s`", response["status"])
break
+ if time.monotonic() - start_time > timeout:
+ raise TimeoutError(
+ f"Query `{query_id}` did not finish within the timeout
period of {timeout} seconds."
+ )
+
time.sleep(poll_interval)
if response["status"] == "error" and raise_error:
diff --git
a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
index eb4fc930c6e..85f93a87ad8 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
@@ -1249,8 +1249,9 @@ class TestSnowflakeSqlApiHook:
assert result == {"status": "queued", "info": ["a", "b"]}
sleep_mock.assert_not_called()
+ @mock.patch(f"{MODULE_PATH}.time.monotonic")
@mock.patch(f"{MODULE_PATH}.time.sleep")
- def test_wait_for_query_timeout_error(self, sleep_mock, time_machine):
+ def test_wait_for_query_timeout_error(self, sleep_mock, monotonic_mock):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
# Simulate a query that keeps running and never finishes
@@ -1259,9 +1260,13 @@ class TestSnowflakeSqlApiHook:
qid = "qid-789"
timeout = 3
- # Freeze the clock. Each sleep advances it explicitly so logger
time.time() calls do not skew the timeout.
- time_machine.move_to(0, tick=False)
- sleep_mock.side_effect = lambda seconds: time_machine.shift(seconds +
0.1)
+ clock = [0.0]
+ monotonic_mock.side_effect = lambda: clock[0]
+
+ def advance_clock(seconds):
+ clock[0] += seconds + 0.1
+
+ sleep_mock.side_effect = advance_clock
with pytest.raises(TimeoutError):
hook.wait_for_query(query_id=qid, timeout=timeout, poll_interval=1)
@@ -1272,6 +1277,28 @@ class TestSnowflakeSqlApiHook:
assert hook.get_sql_api_query_status.call_count == 4
hook.get_sql_api_query_status.assert_has_calls([mock.call(query_id=qid)] * 4)
+ @pytest.mark.parametrize("final_status", ["success", "error"])
+ @mock.patch(f"{MODULE_PATH}.time.monotonic")
+ @mock.patch(f"{MODULE_PATH}.time.sleep")
+ def test_wait_for_query_returns_finished_status_after_timeout_elapsed(
+ self, sleep_mock, monotonic_mock, final_status
+ ):
+ """A status check that outlasts the timeout still reports a query that
has finished."""
+ hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
+ clock = [0.0]
+ monotonic_mock.side_effect = lambda: clock[0]
+
+ def finish_after_slow_status_call(query_id):
+ clock[0] += 10
+ return {"status": final_status, "message": "done"}
+
+ hook.get_sql_api_query_status =
mock.MagicMock(side_effect=finish_after_slow_status_call)
+
+ result = hook.wait_for_query(query_id="qid-slow", timeout=3,
poll_interval=1)
+
+ assert result == {"status": final_status, "message": "done"}
+ sleep_mock.assert_not_called()
+
@mock.patch(f"{HOOK_PATH}._make_api_call_with_retries")
@mock.patch(f"{HOOK_PATH}._process_response")
@mock.patch(f"{HOOK_PATH}.get_request_url_header_params")
