From c7cef500189f08b1b4834faa0ff0132567ce319c Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 6 Jul 2026 14:43:30 +0530 Subject: [PATCH 1/3] Fix SnowflakeSqlApiOperator polling to avoid needless sleeps --- .../snowflake/operators/snowflake.py | 8 ++- .../snowflake/operators/test_snowflake.py | 53 +++++++++++++++++-- 2 files changed, 55 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 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..84684d4b672b9 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -359,6 +359,49 @@ 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_execute_query, 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_execute_query, 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 +624,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 From 092191c08e48d6441991468066de77df28d452f0 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 6 Jul 2026 19:11:26 +0530 Subject: [PATCH 2/3] Add durable execution to SnowflakeSqlApiOperator --- .../snowflake/docs/operators/snowflake.rst | 42 +++ .../snowflake/operators/snowflake.py | 121 ++++++-- .../snowflake/operators/test_snowflake.py | 293 ++++++++++++++++-- 3 files changed, 401 insertions(+), 55 deletions(-) diff --git a/providers/snowflake/docs/operators/snowflake.rst b/providers/snowflake/docs/operators/snowflake.rst index b48930c85a200..85c1df3b2eedc 100644 --- a/providers/snowflake/docs/operators/snowflake.rst +++ b/providers/snowflake/docs/operators/snowflake.rst @@ -154,3 +154,45 @@ An example usage of the SnowflakeSqlApiHook is as follows: Parameters that can be passed onto the operator will be given priority over the parameters already given in the Airflow connection metadata (such as ``schema``, ``role``, ``database`` and so forth). + +Durable execution +^^^^^^^^^^^^^^^^^^ + +``SnowflakeSqlApiOperator`` submits one or more SQL statements and then polls their statement +handles to completion on the worker. By default the operator runs in a *durable* mode that makes +this crash-safe: the statement handles are persisted to task state store before polling begins, so +if the worker crashes or is preempted and the task is retried, the operator reconnects to the +statements that are already executing in Snowflake instead of resubmitting the SQL. + +On retry the operator checks the prior statements' state: + +* if any handle is still running, the operator reconnects and continues polling -- handles that + already finished are not re-run, only the ones still in progress are waited on +* if every handle already succeeded, the operator returns immediately without resubmitting +* if any handle failed, or a handle has expired past Snowflake's retention window, the operator + submits the SQL fresh + +Because Snowflake's SQL API has no way to retry or repair a single failed statement within a +multi-statement request, a genuine failure always resubmits the whole request batch, matching the +operator's all-or-nothing submission semantics. + +Durable execution requires Airflow 3.3 or newer, since it relies on the task state store. On +earlier Airflow versions the flag is a no-op and the operator always submits fresh SQL on retry, +exactly as before. If the task state store is unavailable at runtime, the operator logs that crash +recovery is disabled and behaves the same way. + +To opt out and always submit fresh SQL on retry, set ``durable=False``: + +.. code-block:: python + + api_operator = SnowflakeSqlApiOperator( + task_id="snowflake_sql_api", + snowflake_conn_id="snowflake_default", + sql="select * from table", + statement_count=1, + durable=False, + ) + +Durable execution applies to the synchronous path. When ``deferrable=True`` is set, the Triggerer +already tracks the statement handles across the wait, so deferrable mode takes precedence and +``durable`` has no effect. diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index b9e74b185890a..977906b759f65 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -23,6 +23,8 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, SupportsAbs, cast +import requests + from airflow.providers.common.compat.sdk import conf from airflow.providers.common.sql.operators.sql import ( SQLCheckOperator, @@ -33,6 +35,25 @@ from airflow.providers.snowflake.hooks.snowflake_sql_api import SnowflakeSqlApiHook from airflow.providers.snowflake.triggers.snowflake_trigger import SnowflakeSqlApiTrigger +try: + from airflow.sdk import ResumableJobMixin +except ImportError: + + class ResumableJobMixin: # type: ignore[no-redef] + """Airflow <3.3 stub, task_state_store unavailable, always submits fresh.""" + + external_id_key: str = "snowflake_query_ids" + + def __init__(self, *, durable: bool = True, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.durable = durable + + def execute_resumable(self, context): + external_id = self.submit_job(context) + self.poll_until_complete(external_id, context) + return self.get_job_result(external_id, context) + + if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context @@ -284,7 +305,7 @@ def __init__( self.query_ids: list[str] = [] -class SnowflakeSqlApiOperator(SQLExecuteQueryOperator): +class SnowflakeSqlApiOperator(ResumableJobMixin, SQLExecuteQueryOperator): """ Implemented Snowflake SQL API Operator to support multiple SQL statements sequentially. @@ -357,10 +378,15 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator): To set the timeout to the maximum value (604800 seconds), set timeout to 0. :param deferrable: Run operator in the deferrable mode. :param snowflake_api_retry_args: An optional dictionary with arguments passed to ``tenacity.Retrying`` & ``tenacity.AsyncRetrying`` classes. + :param durable: When ``True`` (the default), the submitted statement handles are persisted to + task state before polling begins. A worker crash on retry reconnects to the existing + statements instead of resubmitting the SQL. Set to ``False`` to always submit fresh on + retry. Requires Airflow 3.3+; ignored silently on earlier versions. """ LIFETIME = timedelta(minutes=59) # The tokens will have a 59 minutes lifetime RENEWAL_DELTA = timedelta(minutes=54) # Tokens will be renewed after 54 minutes + external_id_key = "snowflake_query_ids" template_fields: Sequence[str] = tuple( set(SQLExecuteQueryOperator.template_fields) | {"snowflake_conn_id"} @@ -428,6 +454,9 @@ def execute(self, context: Context) -> None: By deferring the SnowflakeSqlApiTrigger class passed along with query ids. """ + if not self.deferrable: + return self.execute_resumable(context) + self.log.info("Executing: %s", self.sql) self.query_ids = self._hook.execute_query( self.sql, statement_count=self.statement_count, bindings=self.bindings, timeout=self.timeout @@ -452,27 +481,17 @@ def execute(self, context: Context) -> None: self.log.info("%s completed successfully.", self.task_id) return - if self.deferrable: - self.defer( - timeout=self.execution_timeout, - trigger=SnowflakeSqlApiTrigger( - poll_interval=self.poll_interval, - query_ids=self.query_ids, - snowflake_conn_id=self.snowflake_conn_id, - token_life_time=self.token_life_time, - token_renewal_delta=self.token_renewal_delta, - ), - method_name="execute_complete", - ) - else: - while True: - statement_status = self.poll_on_queries() - if statement_status["error"]: - raise RuntimeError(str(statement_status["error"])) - if not statement_status["running"]: - break - - self._hook.check_query_output(self.query_ids) + self.defer( + timeout=self.execution_timeout, + trigger=SnowflakeSqlApiTrigger( + poll_interval=self.poll_interval, + query_ids=self.query_ids, + snowflake_conn_id=self.snowflake_conn_id, + token_life_time=self.token_life_time, + token_renewal_delta=self.token_renewal_delta, + ), + method_name="execute_complete", + ) def poll_on_queries(self): """Poll on requested queries.""" @@ -507,6 +526,64 @@ def poll_on_queries(self): "running": statement_running_status, } + def submit_job(self, context: Context) -> list[str]: + """Submit the SQL for execution and return the resulting statement handles.""" + self.log.info("Executing: %s", self.sql) + self.query_ids = self._hook.execute_query( + self.sql, statement_count=self.statement_count, bindings=self.bindings, timeout=self.timeout + ) + self.log.info("List of query ids %s", self.query_ids) + return self.query_ids + + def get_job_status(self, external_id: list[str], context: Context) -> str: + """Aggregate the status of every handle into a single verdict for the mixin.""" + statuses = [] + for query_id in external_id: + try: + statuses.append(self._hook.get_sql_api_query_status(query_id)["status"]) + except requests.exceptions.HTTPError as e: + if e.response is not None and e.response.status_code == 404: + return "not_found" + raise + if "error" in statuses: + return "error" + if "running" in statuses: + return "running" + return "success" + + def is_job_active(self, status: str) -> bool: + return status == "running" + + def is_job_succeeded(self, status: str) -> bool: + return status == "success" + + def poll_until_complete(self, external_id: list[str], context: Context) -> None: + self.query_ids = external_id + while True: + statement_status = self.poll_on_queries() + if statement_status["error"]: + raise RuntimeError(str(statement_status["error"])) + if not statement_status["running"]: + break + if self.do_xcom_push and context is not None: + context["ti"].xcom_push(key="query_ids", value=self.query_ids) + # On reconnect, the mixin calls poll_until_complete alone -- get_job_result is never + # invoked in that case -- so the output must be fetched here too, not left to + # get_job_result. Fresh submit calls both; the flag stops get_job_result from + # fetching (and pushing xcoms) a second time. + self._hook.check_query_output(self.query_ids) + self._poll_until_complete_ran = True + + def get_job_result(self, external_id: list[str], context: Context) -> None: + self.query_ids = external_id + if getattr(self, "_poll_until_complete_ran", False): + return + # The already-succeeded retry path skips submit_job and poll_until_complete entirely, + # so push the query_ids xcom and fetch output here for parity with the normal path. + if self.do_xcom_push and context is not None: + context["ti"].xcom_push(key="query_ids", value=self.query_ids) + self._hook.check_query_output(self.query_ids) + def execute_complete(self, context: Context, event: dict[str, str | list[str]] | None = None) -> None: """ Execute callback when the trigger fires; returns immediately. diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py index 84684d4b672b9..0b7e707234448 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py @@ -18,10 +18,11 @@ from __future__ import annotations from unittest import mock -from unittest.mock import call +from unittest.mock import MagicMock, call import pendulum import pytest +import requests from airflow.models import Connection from airflow.models.dag import DAG @@ -41,7 +42,7 @@ from tests_common.test_utils.dag import sync_dag_to_db from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_runs from tests_common.test_utils.taskinstance import create_task_instance -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, timezone +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_3_PLUS, timezone DEFAULT_DATE = timezone.datetime(2015, 1, 1) DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat() @@ -272,6 +273,30 @@ def create_context(task, dag=None): } +@pytest.fixture +def mock_execute_query(): + with mock.patch( + "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.execute_query" + ) as execute_query: + yield execute_query + + +@pytest.fixture +def mock_get_sql_api_query_status(): + with mock.patch( + "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.get_sql_api_query_status" + ) as get_sql_api_query_status: + yield get_sql_api_query_status + + +@pytest.fixture +def mock_check_query_output(): + with mock.patch( + "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.check_query_output" + ) as check_query_output: + yield check_query_output + + @pytest.mark.db_test class TestSnowflakeSqlApiOperator: @pytest.fixture(autouse=True) @@ -288,27 +313,6 @@ def setup_tests(self): if AIRFLOW_V_3_0_PLUS: clear_db_dag_bundles() - @pytest.fixture - def mock_execute_query(self): - with mock.patch( - "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.execute_query" - ) as execute_query: - yield execute_query - - @pytest.fixture - def mock_get_sql_api_query_status(self): - with mock.patch( - "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.get_sql_api_query_status" - ) as get_sql_api_query_status: - yield get_sql_api_query_status - - @pytest.fixture - def mock_check_query_output(self): - with mock.patch( - "airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.check_query_output" - ) as check_query_output: - yield check_query_output - def test_snowflake_sql_api_to_succeed_when_no_query_fails( self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output ): @@ -320,11 +324,31 @@ def test_snowflake_sql_api_to_succeed_when_no_query_fails( sql=SQL_MULTIPLE_STMTS, statement_count=4, do_xcom_push=False, + durable=False, ) mock_execute_query.return_value = ["uuid1", "uuid2"] mock_get_sql_api_query_status.side_effect = [{"status": "success"}, {"status": "success"}] operator.execute(context=None) + def test_snowflake_sql_api_durable_true_submits_fresh_missing_task_state_store( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + """Default durable=True with no task_state_store in context still submits and succeeds.""" + operator = SnowflakeSqlApiOperator( + task_id=TASK_ID, + snowflake_conn_id="snowflake_default", + sql=SQL_MULTIPLE_STMTS, + statement_count=4, + do_xcom_push=False, + ) + mock_execute_query.return_value = ["uuid1", "uuid2"] + mock_get_sql_api_query_status.side_effect = [{"status": "success"}, {"status": "success"}] + + operator.execute(context={}) + + mock_execute_query.assert_called_once() + mock_check_query_output.assert_called_once_with(["uuid1", "uuid2"]) + def test_snowflake_sql_api_to_fails_when_one_query_fails( self, mock_execute_query, mock_get_sql_api_query_status ): @@ -336,6 +360,7 @@ def test_snowflake_sql_api_to_fails_when_one_query_fails( sql=SQL_MULTIPLE_STMTS, statement_count=4, do_xcom_push=False, + durable=False, ) mock_execute_query.return_value = ["uuid1", "uuid2"] mock_get_sql_api_query_status.side_effect = [{"status": "error"}, {"status": "success"}] @@ -436,6 +461,7 @@ def test_snowflake_sql_api_pushes_query_ids_to_xcom( self, mock_execute_query, mock_get_sql_api_query_status, + mock_check_query_output, ): """ Tests that query IDs returned by the Snowflake SQL API are pushed to XCom @@ -448,6 +474,7 @@ def test_snowflake_sql_api_pushes_query_ids_to_xcom( statement_count=4, do_xcom_push=True, deferrable=False, + durable=False, ) mock_execute_query.return_value = ["uuid1"] @@ -617,27 +644,29 @@ def test_snowflake_sql_api_execute_operator_polling_running( statement_count=4, do_xcom_push=False, deferrable=False, + durable=False, ) mock_execute_query.return_value = ["uuid1"] 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) -- still running, sleeps {"status": "running"}, # 2nd poll_on_queries check (poll_interval: 5s) -- still running, sleeps {"status": "running"}, - # 3rd poll_on_queries check -- resolves to success, no sleep needed + # 3rd poll_on_queries check (poll_interval: 5s) -- still running, sleeps + {"status": "running"}, + # 4th 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"]) - # 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 + # 3 sleeps: durable=False routes straight into poll_until_complete with no + # separate pre-check, so every "running" cycle (including the first) sleeps; + # only the cycle that resolves to success skips it. + assert mock_sleep.call_count == 3 def test_snowflake_sql_api_execute_operator_polling_failed( self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output @@ -653,19 +682,24 @@ def test_snowflake_sql_api_execute_operator_polling_failed( statement_count=4, do_xcom_push=False, deferrable=False, + durable=False, ) mock_execute_query.return_value = ["uuid1"] mock_get_sql_api_query_status.side_effect = [ - # Initial get_sql_api_query_status check + # 1st poll_on_queries check -- still running, sleeps {"status": "running"}, - # 1st poll_on_queries check + # 2nd poll_on_queries check -- resolves to error, raises immediately {"status": "error"}, ] - with pytest.raises(RuntimeError): - operator.execute(context=None) + with mock.patch("time.sleep") as mock_sleep: + with pytest.raises(RuntimeError): + operator.execute(context=None) + # 1 sleep: the first cycle finds "running" and sleeps before the second + # cycle finds "error" and raises without sleeping again. + assert mock_sleep.call_count == 1 mock_check_query_output.assert_not_called() @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries") @@ -697,3 +731,196 @@ def test_snowflake_sql_api_on_kill_no_queries(self, mock_cancel_queries): operator.on_kill() mock_cancel_queries.assert_not_called() + + +@pytest.mark.skipif( + not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution) requires Airflow 3.3+" +) +class TestSnowflakeSqlApiOperatorDurable: + @staticmethod + def _context(task_store=None): + ctx: dict = {"ti": MagicMock(stats_tags={})} + if task_store is not None: + ctx["task_state_store"] = task_store + return ctx + + @staticmethod + def _make_operator(**kwargs): + return SnowflakeSqlApiOperator( + task_id=TASK_ID, + snowflake_conn_id="snowflake_default", + sql=SQL_MULTIPLE_STMTS, + statement_count=4, + do_xcom_push=False, + **kwargs, + ) + + def test_persists_query_ids_to_task_state_store_on_fresh_submit( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + mock_execute_query.return_value = ["uuid1"] + mock_get_sql_api_query_status.return_value = {"status": "success"} + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = None + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_called_once() + task_store.set.assert_called_once_with("snowflake_query_ids", ["uuid1"]) + + def test_reconnects_to_running_query_without_resubmitting( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + # get_job_status sees it still running, then poll_on_queries sees it finish. + mock_get_sql_api_query_status.side_effect = [{"status": "running"}, {"status": "success"}] + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = ["uuid1"] + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_not_called() + task_store.set.assert_not_called() + + def test_partial_progress_reconnect_waits_only_on_running_handle( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + mock_get_sql_api_query_status.side_effect = [ + {"status": "success"}, # get_job_status check: uuid1 already done + {"status": "running"}, # get_job_status check: uuid2 -> aggregate "running" -> reconnect + {"status": "success"}, # poll cycle 1: uuid1 re-confirmed + {"status": "running"}, # poll cycle 1: uuid2 still running -> sleep + {"status": "success"}, # poll cycle 2: uuid1 re-confirmed + {"status": "success"}, # poll cycle 2: uuid2 finishes + ] + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = ["uuid1", "uuid2"] + + with mock.patch("time.sleep"): + operator.execute(self._context(task_store)) + + mock_execute_query.assert_not_called() + task_store.set.assert_not_called() + mock_check_query_output.assert_called_once_with(["uuid1", "uuid2"]) + + def test_already_succeeded_returns_result_without_polling( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + mock_get_sql_api_query_status.return_value = {"status": "success"} + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = ["uuid1"] + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_not_called() + mock_check_query_output.assert_called_once_with(["uuid1"]) + # Only the get_job_status check ran; poll_on_queries must never have been entered. + assert mock_get_sql_api_query_status.call_count == 1 + + def test_resubmits_when_stored_query_in_terminal_error( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + mock_execute_query.return_value = ["uuid2"] + # get_job_status sees the stored handle failed; after fresh resubmit, polling succeeds. + mock_get_sql_api_query_status.side_effect = [{"status": "error"}, {"status": "success"}] + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = ["uuid1"] + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_called_once() + task_store.set.assert_called_once_with("snowflake_query_ids", ["uuid2"]) + + def test_resubmits_when_stored_query_not_found( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator() + mock_execute_query.return_value = ["uuid2"] + not_found_response = MagicMock() + not_found_response.status_code = 404 + # get_job_status sees the stored handle expired (404); after fresh resubmit, polling succeeds. + mock_get_sql_api_query_status.side_effect = [ + requests.exceptions.HTTPError(response=not_found_response), + {"status": "success"}, + ] + task_store = MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = ["uuid1"] + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_called_once() + task_store.set.assert_called_once_with("snowflake_query_ids", ["uuid2"]) + + def test_durable_false_never_touches_task_state_store( + self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output + ): + operator = self._make_operator(durable=False) + mock_execute_query.return_value = ["uuid1"] + mock_get_sql_api_query_status.return_value = {"status": "success"} + task_store = MagicMock(spec_set=["get", "set"]) + + operator.execute(self._context(task_store)) + + mock_execute_query.assert_called_once() + task_store.get.assert_not_called() + task_store.set.assert_not_called() + + def test_deferrable_unaffected_by_durable(self, mock_execute_query, mock_get_sql_api_query_status): + operator = self._make_operator(deferrable=True) + mock_execute_query.return_value = ["uuid1"] + mock_get_sql_api_query_status.return_value = {"status": "running"} + task_store = MagicMock(spec_set=["get", "set"]) + + with mock.patch.object(SnowflakeSqlApiOperator, "defer") as mock_defer: + operator.execute(self._context(task_store)) + + assert mock_defer.called + task_store.get.assert_not_called() + task_store.set.assert_not_called() + + def test_get_job_status_error_takes_priority_over_running(self, mock_get_sql_api_query_status): + operator = self._make_operator() + mock_get_sql_api_query_status.side_effect = [{"status": "running"}, {"status": "error"}] + + assert operator.get_job_status(["uuid1", "uuid2"], context={}) == "error" + + def test_get_job_status_running_when_none_error(self, mock_get_sql_api_query_status): + operator = self._make_operator() + mock_get_sql_api_query_status.side_effect = [{"status": "success"}, {"status": "running"}] + + assert operator.get_job_status(["uuid1", "uuid2"], context={}) == "running" + + def test_get_job_status_success_when_all_succeed(self, mock_get_sql_api_query_status): + operator = self._make_operator() + mock_get_sql_api_query_status.return_value = {"status": "success"} + + assert operator.get_job_status(["uuid1", "uuid2"], context={}) == "success" + + def test_get_job_status_not_found_on_404(self, mock_get_sql_api_query_status): + operator = self._make_operator() + response = MagicMock() + response.status_code = 404 + mock_get_sql_api_query_status.side_effect = requests.exceptions.HTTPError(response=response) + + assert operator.get_job_status(["uuid1"], context={}) == "not_found" + + def test_get_job_status_reraises_non_404_http_error(self, mock_get_sql_api_query_status): + operator = self._make_operator() + response = MagicMock() + response.status_code = 500 + mock_get_sql_api_query_status.side_effect = requests.exceptions.HTTPError(response=response) + + with pytest.raises(requests.exceptions.HTTPError): + operator.get_job_status(["uuid1"], context={}) + + def test_is_job_active_and_is_job_succeeded_predicates(self): + operator = self._make_operator() + + assert operator.is_job_active("running") is True + assert operator.is_job_active("success") is False + assert operator.is_job_succeeded("success") is True + assert operator.is_job_succeeded("running") is False From a02ce7ac7b498e294628d3083d5f0b3a6c76ac44 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 7 Jul 2026 10:25:55 +0530 Subject: [PATCH 3/3] handling mypy --- .../providers/snowflake/operators/snowflake.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index 977906b759f65..a66ff2771e429 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -55,6 +55,8 @@ def execute_resumable(self, context): if TYPE_CHECKING: + from pydantic import JsonValue + from airflow.providers.common.compat.sdk import Context @@ -526,19 +528,19 @@ def poll_on_queries(self): "running": statement_running_status, } - def submit_job(self, context: Context) -> list[str]: + def submit_job(self, context: Context) -> JsonValue: """Submit the SQL for execution and return the resulting statement handles.""" self.log.info("Executing: %s", self.sql) self.query_ids = self._hook.execute_query( self.sql, statement_count=self.statement_count, bindings=self.bindings, timeout=self.timeout ) self.log.info("List of query ids %s", self.query_ids) - return self.query_ids + return cast("JsonValue", self.query_ids) - def get_job_status(self, external_id: list[str], context: Context) -> str: + def get_job_status(self, external_id: JsonValue, context: Context) -> str: """Aggregate the status of every handle into a single verdict for the mixin.""" statuses = [] - for query_id in external_id: + for query_id in cast("list[str]", external_id): try: statuses.append(self._hook.get_sql_api_query_status(query_id)["status"]) except requests.exceptions.HTTPError as e: @@ -557,8 +559,8 @@ def is_job_active(self, status: str) -> bool: def is_job_succeeded(self, status: str) -> bool: return status == "success" - def poll_until_complete(self, external_id: list[str], context: Context) -> None: - self.query_ids = external_id + def poll_until_complete(self, external_id: JsonValue, context: Context) -> None: + self.query_ids = cast("list[str]", external_id) while True: statement_status = self.poll_on_queries() if statement_status["error"]: @@ -574,8 +576,8 @@ def poll_until_complete(self, external_id: list[str], context: Context) -> None: self._hook.check_query_output(self.query_ids) self._poll_until_complete_ran = True - def get_job_result(self, external_id: list[str], context: Context) -> None: - self.query_ids = external_id + def get_job_result(self, external_id: JsonValue, context: Context) -> None: + self.query_ids = cast("list[str]", external_id) if getattr(self, "_poll_until_complete_ran", False): return # The already-succeeded retry path skips submit_job and poll_until_complete entirely,