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 96b820cad5d Add durable execution to SnowflakeSqlApiOperator (#69477)
96b820cad5d is described below
commit 96b820cad5da43387bd1ad404cd5c0c33849eda9
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Jul 14 11:25:58 2026 +0530
Add durable execution to SnowflakeSqlApiOperator (#69477)
---
providers/snowflake/docs/operators/snowflake.rst | 42 +++
.../providers/snowflake/operators/snowflake.py | 132 ++++++--
.../unit/snowflake/operators/test_snowflake.py | 342 ++++++++++++++++++---
3 files changed, 459 insertions(+), 57 deletions(-)
diff --git a/providers/snowflake/docs/operators/snowflake.rst
b/providers/snowflake/docs/operators/snowflake.rst
index b48930c85a2..85c1df3b2ee 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 b9e74b18589..8e3cda63ef3 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 datetime import timedelta
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,7 +35,28 @@ from airflow.providers.common.sql.operators.sql import (
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 pydantic import JsonValue
+
from airflow.providers.common.compat.sdk import Context
@@ -284,7 +307,7 @@ class
SnowflakeIntervalCheckOperator(SQLIntervalCheckOperator):
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 +380,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 +456,9 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
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 +483,17 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
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 +528,73 @@ class SnowflakeSqlApiOperator(SQLExecuteQueryOperator):
"running": statement_running_status,
}
+ 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 cast("JsonValue", self.query_ids)
+
+ 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 cast("list[str]", 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: JsonValue, context: Context) ->
None:
+ self.query_ids = cast("list[str]", external_id)
+ # On reconnect, execute_query (the only thing that normally populates
this) never ran
+ # on this hook instance -- sync it so OpenLineage's
get_openlineage_database_specific_lineage,
+ # which reads hook.query_ids (not the operator's), doesn't silently
produce no lineage.
+ self._hook.query_ids = self.query_ids
+ # Push before polling, not after, so the handles are recorded even if
a statement
+ # errors below.
+ if self.do_xcom_push and context is not None:
+ context["ti"].xcom_push(key="query_ids", value=self.query_ids)
+ while True:
+ statement_status = self.poll_on_queries()
+ if statement_status["error"]:
+ raise RuntimeError(str(statement_status["error"]))
+ if not statement_status["running"]:
+ break
+ # 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: JsonValue, context: Context) -> None:
+ self.query_ids = cast("list[str]", external_id)
+ # Same reconnect-hook gap as poll_until_complete -- see the comment
there. This path
+ # hits it too, since the already-succeeded case never calls
poll_until_complete either.
+ self._hook.query_ids = self.query_ids
+ 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 b1d6b418839..d37f420f2a6 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 airflow.utils.types import DagRunType
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):
}
[email protected]
+def mock_execute_query():
+ with mock.patch(
+
"airflow.providers.snowflake.operators.snowflake.SnowflakeSqlApiHook.execute_query"
+ ) as execute_query:
+ yield execute_query
+
+
[email protected]
+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
+
+
[email protected]
+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 @@ class TestSnowflakeSqlApiOperator:
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 @@ class TestSnowflakeSqlApiOperator:
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 @@ class TestSnowflakeSqlApiOperator:
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"}]
@@ -359,7 +384,9 @@ 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):
+ 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",
@@ -378,7 +405,7 @@ class TestSnowflakeSqlApiOperator:
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):
+ 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,
@@ -434,6 +461,7 @@ class TestSnowflakeSqlApiOperator:
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
@@ -446,6 +474,7 @@ class TestSnowflakeSqlApiOperator:
statement_count=4,
do_xcom_push=True,
deferrable=False,
+ durable=False,
)
mock_execute_query.return_value = ["uuid1"]
@@ -615,27 +644,29 @@ class TestSnowflakeSqlApiOperator:
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
@@ -651,21 +682,44 @@ class TestSnowflakeSqlApiOperator:
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()
+ def test_poll_until_complete_pushes_query_ids_to_xcom_even_on_failure(
+ self, mock_get_sql_api_query_status
+ ):
+ operator = SnowflakeSqlApiOperator(
+ task_id=TASK_ID,
+ snowflake_conn_id=CONN_ID,
+ sql=SQL_MULTIPLE_STMTS,
+ statement_count=4,
+ do_xcom_push=True,
+ )
+ mock_get_sql_api_query_status.side_effect = [{"status": "error"}]
+ context = create_context(operator)
+
+ with pytest.raises(RuntimeError):
+ operator.poll_until_complete(["uuid1"], context)
+
+ context["ti"].xcom_push.assert_called_once_with(key="query_ids",
value=["uuid1"])
+
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries")
def test_snowflake_sql_api_on_kill_cancels_queries(self,
mock_cancel_queries):
"""Test that on_kill cancels running queries."""
@@ -695,3 +749,221 @@ class TestSnowflakeSqlApiOperator:
operator.on_kill()
mock_cancel_queries.assert_not_called()
+
+
[email protected](
+ 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()
+ # execute_query never ran on this hook instance, so nothing else would
populate this --
+ # OpenLineage's get_openlineage_database_specific_lineage reads
hook.query_ids directly.
+ assert operator._hook.query_ids == ["uuid1"]
+
+ 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"])
+ assert operator._hook.query_ids == ["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
+ assert operator._hook.query_ids == ["uuid1"]
+
+ def test_already_succeeded_pushes_query_ids_to_xcom(
+ self, mock_execute_query, mock_get_sql_api_query_status,
mock_check_query_output
+ ):
+ operator = SnowflakeSqlApiOperator(
+ task_id=TASK_ID,
+ snowflake_conn_id="snowflake_default",
+ sql=SQL_MULTIPLE_STMTS,
+ statement_count=4,
+ do_xcom_push=True,
+ )
+ mock_get_sql_api_query_status.return_value = {"status": "success"}
+ task_store = MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = ["uuid1"]
+ context = self._context(task_store)
+
+ operator.execute(context)
+
+ mock_execute_query.assert_not_called()
+ context["ti"].xcom_push.assert_called_once_with(key="query_ids",
value=["uuid1"])
+
+ 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
