diff --git a/providers/snowflake/docs/operators/snowpark_containers.rst b/providers/snowflake/docs/operators/snowpark_containers.rst index 3d475350db134..b28b8ad873677 100644 --- a/providers/snowflake/docs/operators/snowpark_containers.rst +++ b/providers/snowflake/docs/operators/snowpark_containers.rst @@ -75,3 +75,29 @@ An example usage of the SnowparkContainerJobOperator 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 +^^^^^^^^^^^^^^^^^ + +By default the operator runs in a *durable* mode that makes the synchronous poll crash-safe. +Before polling begins the job name is persisted to :doc:`task state store +`, so if the worker crashes or is preempted and the +task is retried, the operator reconnects to the job already running in Snowflake instead of +submitting a duplicate. + +On retry the operator checks the prior job's status: + +* still running: reconnect and continue polling +* already succeeded: return immediately without resubmitting +* failed or another terminal state: submit a fresh job + +Durable execution requires Airflow 3.3 or newer, since it relies on the task state store. On +earlier versions the flag is a no-op (setting it only emits a warning) and the operator always +submits a fresh job on retry. Durable execution only applies to the synchronous polling path. It has no effect when +``wait_for_completion=False`` or ``deferrable=True``. Set ``durable=False`` to opt out and always +submit a fresh job on retry. + +.. note:: + + With a fixed ``name``, retries after a terminal failure can hit an "already exists" error, since + failed services are kept. Omit ``name`` to auto-generate a unique one if jobs may fail and retry. diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py index 3312e35019a20..d0d0cc7396f54 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py @@ -18,10 +18,13 @@ from __future__ import annotations import time +import warnings from collections.abc import Sequence from datetime import timedelta from functools import cached_property -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast + +from snowflake.connector.errors import ProgrammingError from airflow.providers.common.compat.sdk import conf from airflow.providers.common.compat.standard.operators import BaseOperator @@ -30,15 +33,54 @@ from airflow.providers.snowflake.triggers.snowpark_containers import SnowparkContainerJobTrigger from airflow.providers.snowflake.utils.snowpark_containers import ( NON_TERMINAL_STATUSES, + NOT_FOUND_STATUS, + OBJECT_NOT_EXIST_ERROR_CODE, TERMINAL_STATUSES, SnowparkContainerJobStatus, ) +_DURABLE_UNSET = object() + + +def _warn_and_disable_durable_pre_3_3(durable: Any) -> bool: + """Disable durable below 3.3, warning if it was explicitly set.""" + if durable is not _DURABLE_UNSET: + warnings.warn( + "`durable` has no effect on Airflow versions below 3.3.", + UserWarning, + stacklevel=3, + ) + return False + + +# ResumableJobMixin only exists on Airflow 3.3+ and this provider still targets >=2.11. Drop this +# fallback once the provider's minimum Airflow version is >=3.3. +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 = "snowpark_container_job_name" + + def __init__(self, *, durable: Any = _DURABLE_UNSET, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.durable = _warn_and_disable_durable_pre_3_3(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 -class SnowparkContainerJobOperator(BaseOperator): +class SnowparkContainerJobOperator(ResumableJobMixin, BaseOperator): """ Execute a job on Snowpark Container Services. @@ -81,6 +123,11 @@ class SnowparkContainerJobOperator(BaseOperator): ``wait_for_completion`` is True. With ``wait_for_completion=False`` the operator submits the job and returns immediately without deferring. (default value: False) + :param durable: When ``True``, the submitted job name is persisted to + task state before polling begins. A worker crash on retry reconnects to the existing + job instead of resubmitting the SQL. Set to ``False`` to always submit fresh on + retry. Requires Airflow 3.3+; ignored on earlier versions. + With ``wait_for_completion=False`` or ``deferrable=True`` durable has no effect. (default value: True) :param timeout: Maximum seconds to wait for the job to reach a terminal state. When it elapses the task fails. (default value: 86400) :param database: name of database (will overwrite database defined @@ -94,6 +141,7 @@ class SnowparkContainerJobOperator(BaseOperator): own SQL commands, not for the container's queries """ + external_id_key = "snowpark_container_job_name" template_fields: Sequence[str] = ( "compute_pool", "spec", @@ -123,6 +171,7 @@ def __init__( poll_interval: int = 10, snowflake_conn_id: str = "snowflake_default", deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + durable: bool | None = None, timeout: int = 24 * 60 * 60, database: str | None = None, schema: str | None = None, @@ -130,9 +179,13 @@ def __init__( warehouse: str | None = None, **kwargs: Any, ) -> None: - super().__init__(**kwargs) if spec_text is not None and (spec is not None or spec_stage is not None): raise ValueError("Cannot specify both 'spec_text' and 'spec'/'spec_stage'") + # durable is a named parameter here (not left to **kwargs) so default_args={"durable": ...} + # reaches it on every supported Airflow version. + if durable is not None: + kwargs["durable"] = durable + super().__init__(**kwargs) self.compute_pool = compute_pool self.container_name = container_name self.spec = spec @@ -154,6 +207,9 @@ def __init__( self.warehouse = warehouse # Set after the job is submitted, parsed from the job submission response. self.job_name: str | None = None + # On a fresh submit the mixin runs both poll_until_complete and get_job_result. + # poll_until_complete sets this so get_job_result does not finalize a second time. + self._poll_until_complete_ran = False if self.deferrable and not self.wait_for_completion: self.log.warning("deferrable has no effect when wait_for_completion is False.") @@ -191,14 +247,41 @@ def _run_one(self, sql: str, return_dictionaries: bool = False) -> Any: """Run a single statement that returns one row via fetch_one_handler.""" return self._hook.run(sql, handler=fetch_one_handler, return_dictionaries=return_dictionaries) - def _submit_job(self) -> str: + def _describe_status(self, external_id: JsonValue) -> str: + """Describe the job's current status.""" + response = self._run_one(f"DESCRIBE SERVICE {external_id}", return_dictionaries=True) + return response.get("status") + + def submit_job(self, context: Context) -> str: """Submit the job and return the name.""" response = self._run_one(self._build_sql()) - job_name = response[0].split("'")[1] - return job_name + self.job_name = response[0].split("'")[1] + if not self.job_name: + raise RuntimeError("Job name was not returned") + return self.job_name + + def get_job_status(self, external_id: JsonValue, context: Context) -> str: + """Return NOT_FOUND when the service no longer exists, otherwise the current status.""" + try: + return self._describe_status(external_id) + except ProgrammingError as e: + if e.errno == OBJECT_NOT_EXIST_ERROR_CODE: + return NOT_FOUND_STATUS + raise + + def is_job_active(self, status: str) -> bool: + """Return True while the job is still running.""" + return status in NON_TERMINAL_STATUSES + + def is_job_succeeded(self, status: str) -> bool: + """Return True when the job has completed successfully.""" + return status == SnowparkContainerJobStatus.DONE + + def poll_until_complete(self, external_id: JsonValue, context: Context) -> None: + """Poll the job until it reaches a terminal state and handle the final status.""" + # On reconnect the mixin skips submit_job, so set the job name from the external id here. + self.job_name = cast("str", external_id) - def _poll_for_status(self) -> str: - """Poll until the job reaches a terminal state.""" status = None end_time = time.monotonic() + self.timeout while True: @@ -207,14 +290,25 @@ def _poll_for_status(self) -> str: if self.drop_on_completion: self._drop_service() raise TimeoutError(f"Job {self.job_name} did not reach a terminal status before the timeout.") - response = self._run_one(f"DESCRIBE SERVICE {self.job_name}", return_dictionaries=True) - status = response.get("status") + status = self._describe_status(self.job_name) if status in TERMINAL_STATUSES: - return status + # get_job_result is skipped when the mixin reconnects to a still-running job, so + # finalize here. + self._handle_final_status(status=status) + self._poll_until_complete_ran = True + return if status not in NON_TERMINAL_STATUSES: raise RuntimeError(f"Job {self.job_name} returned unexpected status: {status}") time.sleep(self.poll_interval) + def get_job_result(self, external_id: JsonValue, context: Context) -> None: + """Finalize the completed job unless poll_until_complete already did.""" + self.job_name = cast("str", external_id) + if self._poll_until_complete_ran: + return + # The mixin only reaches this path when the job is DONE, so the status is hardcoded. + self._handle_final_status(status=SnowparkContainerJobStatus.DONE) + def _log_container_output(self, status: str | None) -> None: """Fetch and log container output for all replicas. Best-effort so it never blocks cleanup.""" for instance_id in range(self.replicas): @@ -259,12 +353,11 @@ def execute(self, context: Context) -> str: if not self.spec_text and not (self.spec and self.spec_stage): raise ValueError("Must provide either 'spec_text' or both 'spec' and 'spec_stage'") - self.job_name = self._submit_job() - if not self.job_name: - raise RuntimeError("Job name was not returned") if not self.wait_for_completion: - return self.job_name + return self.submit_job(context) + if self.deferrable: + job_name = self.submit_job(context) # timeout and execution_timeout give the trigger two separate deadlines. timeout caps # how long the job is polled, and execution_timeout, when set, enforces the task-level # limit. The trigger times out on whichever is reached first. @@ -282,7 +375,7 @@ def execute(self, context: Context) -> str: defer_timeout = self.execution_timeout + poll_buffer self.defer( trigger=SnowparkContainerJobTrigger( - job_name=self.job_name, + job_name=job_name, snowflake_conn_id=self.snowflake_conn_id, poll_interval=self.poll_interval, end_time=now + self.timeout, @@ -295,9 +388,8 @@ def execute(self, context: Context) -> str: timeout=defer_timeout, method_name="execute_complete", ) - status = self._poll_for_status() - self._handle_final_status(status) - return self.job_name + self.execute_resumable(context) + return cast("str", self.job_name) def execute_complete(self, context: Context, event: dict[str, Any]) -> str: """Resume after the trigger fires.""" diff --git a/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py b/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py index 001446123b1a9..52719c85e3293 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py +++ b/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py @@ -18,6 +18,9 @@ from enum import Enum +NOT_FOUND_STATUS = "NOT_FOUND" +OBJECT_NOT_EXIST_ERROR_CODE = 2003 + class SnowparkContainerJobStatus(str, Enum): """Statuses of a Snowpark Container Services job service.""" diff --git a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py index fa888c383f809..851ffb5221edf 100644 --- a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py +++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py @@ -17,14 +17,26 @@ from __future__ import annotations import itertools +import warnings from datetime import datetime, timedelta, timezone from unittest import mock import pytest +from snowflake.connector.errors import ProgrammingError from airflow.providers.common.compat.sdk import TaskDeferred -from airflow.providers.snowflake.operators.snowpark_containers import SnowparkContainerJobOperator +from airflow.providers.snowflake.operators.snowpark_containers import ( + _DURABLE_UNSET, + SnowparkContainerJobOperator, + _warn_and_disable_durable_pre_3_3, +) from airflow.providers.snowflake.triggers.snowpark_containers import SnowparkContainerJobTrigger +from airflow.providers.snowflake.utils.snowpark_containers import ( + NOT_FOUND_STATUS, + OBJECT_NOT_EXIST_ERROR_CODE, +) + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS TASK_ID = "test_spcs_job" COMPUTE_POOL = "test_pool" @@ -35,6 +47,7 @@ JOB_NAME = "TEST_JOB" SNOWFLAKE_CONN_ID = "snowflake_default" MOCK_HOOK_PATH = "airflow.providers.snowflake.operators.snowpark_containers.SnowflakeHook" +SUBMIT_RESPONSE = [f"Started Snowpark Container Services Job '{JOB_NAME}'."] def _make_operator(**kwargs): @@ -49,6 +62,13 @@ def _make_operator(**kwargs): return SnowparkContainerJobOperator(**defaults) +def _context(task_store=None): + ctx = {"ti": mock.MagicMock(stats_tags={})} + if task_store is not None: + ctx["task_state_store"] = task_store + return ctx + + class TestSnowparkContainerJobOperator: def test_invalid_spec_combinations_at_init(self): with pytest.raises(ValueError, match=r"Cannot specify both"): @@ -141,38 +161,42 @@ def test_external_access_integrations_in_template_fields(self): assert "external_access_integrations" in op.template_fields assert hasattr(op, "external_access_integrations") + def test_external_id_key(self): + assert _make_operator().external_id_key == "snowpark_container_job_name" + @mock.patch(MOCK_HOOK_PATH) def test_submit_job_parses_job_name(self, mock_hook_cls): mock_hook = mock_hook_cls.return_value - mock_hook.run.return_value = ["Started Snowpark Container Services Job 'TEST_JOB'."] + mock_hook.run.return_value = SUBMIT_RESPONSE op = _make_operator() - result = op._submit_job() + result = op.submit_job(None) assert result == JOB_NAME @pytest.mark.parametrize( "status", - ("DONE", "FAILED", "CANCELLED", "INTERNAL_ERROR"), + ("FAILED", "CANCELLED", "INTERNAL_ERROR"), ) + @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") @mock.patch(MOCK_HOOK_PATH) - def test_poll_returns_terminal_status(self, mock_hook_cls, status): + def test_poll_raises_on_terminal_failure(self, mock_hook_cls, mock_log, status): mock_hook = mock_hook_cls.return_value mock_hook.run.return_value = {"status": status} op = _make_operator(poll_interval=0) - op.job_name = JOB_NAME - assert op._poll_for_status() == status + with pytest.raises(RuntimeError, match="finished with status"): + op.poll_until_complete(JOB_NAME, None) @mock.patch(MOCK_HOOK_PATH) def test_poll_raises_on_unexpected_status(self, mock_hook_cls): mock_hook = mock_hook_cls.return_value mock_hook.run.return_value = {"status": "UNKNOWN"} op = _make_operator(poll_interval=0) - op.job_name = JOB_NAME with pytest.raises(RuntimeError, match="unexpected status"): - op._poll_for_status() + op.poll_until_complete(JOB_NAME, None) + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") @mock.patch("time.sleep") @mock.patch(MOCK_HOOK_PATH) - def test_poll_waits_through_pending_then_done(self, mock_hook_cls, mock_sleep): + def test_poll_waits_through_pending_then_done(self, mock_hook_cls, mock_sleep, mock_handle): mock_hook = mock_hook_cls.return_value mock_hook.run.side_effect = [ {"status": "PENDING"}, @@ -180,9 +204,9 @@ def test_poll_waits_through_pending_then_done(self, mock_hook_cls, mock_sleep): {"status": "DONE"}, ] op = _make_operator(poll_interval=5) - op.job_name = JOB_NAME - assert op._poll_for_status() == "DONE" + op.poll_until_complete(JOB_NAME, None) assert mock_sleep.call_count == 2 + mock_handle.assert_called_once_with(status="DONE") @pytest.mark.parametrize( ("drop_on_completion", "drops"), @@ -202,7 +226,7 @@ def test_poll_raises_and_logs_on_timeout( op.job_name = JOB_NAME with pytest.raises(TimeoutError, match="did not reach a terminal status"): - op._poll_for_status() + op.poll_until_complete(JOB_NAME, None) mock_log.assert_called_once_with("RUNNING") drop_call = mock.call(f"DROP SERVICE IF EXISTS {JOB_NAME}") @@ -292,66 +316,64 @@ def test_submit_job_raises_on_malformed_response(self, mock_hook_cls): mock_hook.run.return_value = ["unexpected response"] op = _make_operator() with pytest.raises(IndexError): - op._submit_job() + op.submit_job(None) - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=None) - def test_execute_raises_when_job_name_not_returned(self, mock_submit): + @mock.patch(MOCK_HOOK_PATH) + def test_submit_job_raises_when_job_name_empty(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = ["Started Snowpark Container Services Job ''."] op = _make_operator() with pytest.raises(RuntimeError, match="Job name was not returned"): - op.execute(context=None) + op.submit_job(None) @mock.patch(MOCK_HOOK_PATH) def test_execute_no_wait(self, mock_hook_cls): mock_hook = mock_hook_cls.return_value - mock_hook.run.return_value = ["Started Snowpark Container Services Job 'TEST_JOB'."] + mock_hook.run.return_value = SUBMIT_RESPONSE op = _make_operator(wait_for_completion=False) result = op.execute(context=None) assert result == JOB_NAME assert mock_hook.run.call_count == 1 - @mock.patch(MOCK_HOOK_PATH) @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") - @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) - def test_execute_wait_success(self, mock_submit, mock_poll, mock_log, mock_hook_cls): - op = _make_operator() + @mock.patch.object(SnowparkContainerJobOperator, "poll_until_complete") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) + @mock.patch(MOCK_HOOK_PATH) + def test_execute_wait_success(self, mock_hook_cls, mock_submit, mock_poll, mock_log): + op = _make_operator(durable=False) result = op.execute(context=None) mock_submit.assert_called_once() mock_poll.assert_called_once() mock_log.assert_called_once_with("DONE") assert result == JOB_NAME - @mock.patch(MOCK_HOOK_PATH) - @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") - @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="FAILED") - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) - def test_execute_wait_failure_raises(self, mock_submit, mock_poll, mock_log, mock_hook_cls): - op = _make_operator() - with pytest.raises(RuntimeError, match="FAILED"): - op.execute(context=None) - mock_log.assert_called_once_with("FAILED") - @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") - @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) @mock.patch(MOCK_HOOK_PATH) - def test_execute_drops_service_on_completion(self, mock_hook_cls, mock_submit, mock_poll, mock_log): + def test_execute_drops_service_on_completion(self, mock_hook_cls, mock_submit, mock_log): mock_hook = mock_hook_cls.return_value - op = _make_operator(drop_on_completion=True) + mock_hook.run.side_effect = [ + {"status": "DONE"}, + None, + ] + op = _make_operator(drop_on_completion=True, poll_interval=0, durable=False) op.execute(context=None) - mock_hook.run.assert_called_once_with(f"DROP SERVICE IF EXISTS {JOB_NAME}") + assert mock.call(f"DROP SERVICE IF EXISTS {JOB_NAME}") in mock_hook.run.call_args_list @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output") - @mock.patch.object(SnowparkContainerJobOperator, "_poll_for_status", return_value="DONE") - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) @mock.patch(MOCK_HOOK_PATH) - def test_execute_skips_drop_when_disabled(self, mock_hook_cls, mock_submit, mock_poll, mock_log): + def test_execute_skips_drop_when_disabled(self, mock_hook_cls, mock_submit, mock_log): mock_hook = mock_hook_cls.return_value - op = _make_operator(drop_on_completion=False) + mock_hook.run.side_effect = [ + {"status": "DONE"}, + ] + op = _make_operator(drop_on_completion=False, poll_interval=0, durable=False) op.execute(context=None) - mock_hook.run.assert_not_called() + drop_call = mock.call(f"DROP SERVICE IF EXISTS {JOB_NAME}") + assert drop_call not in mock_hook.run.call_args_list - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) def test_execute_defers_when_deferrable(self, mock_submit): op = _make_operator(deferrable=True) with pytest.raises(TaskDeferred) as exc: @@ -360,7 +382,7 @@ def test_execute_defers_when_deferrable(self, mock_submit): assert exc.value.trigger.job_name == JOB_NAME assert exc.value.method_name == "execute_complete" - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) def test_execute_defer_without_execution_timeout(self, mock_submit): op = _make_operator(deferrable=True, timeout=100, poll_interval=10) with pytest.raises(TaskDeferred) as exc: @@ -368,7 +390,7 @@ def test_execute_defer_without_execution_timeout(self, mock_submit): assert exc.value.trigger.execution_deadline is None assert exc.value.timeout == timedelta(seconds=100 + 10 + 60) - @mock.patch.object(SnowparkContainerJobOperator, "_submit_job", return_value=JOB_NAME) + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) def test_execute_defer_uses_execution_timeout_for_deadline_and_buffer(self, mock_submit, time_machine): time_machine.move_to(1000, tick=False) context = {"ti": mock.Mock(start_date=datetime.fromtimestamp(1000, tz=timezone.utc))} @@ -449,3 +471,214 @@ def test_execute_complete_timeout_skips_drop_when_disabled(self, mock_log, mock_ ) mock_log.assert_called_once_with("timeout") mock_hook.run.assert_not_called() + + +@pytest.mark.skipif( + not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution) requires Airflow 3.3+" +) +class TestSnowparkContainerJobOperatorDurable: + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) + @mock.patch(MOCK_HOOK_PATH) + def test_job_name_persists_to_task_state_store_on_fresh_submit( + self, mock_hook_cls, mock_submit_job, mock_handle + ): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = {"status": "DONE"} + + op = _make_operator(poll_interval=0) + task_store = mock.MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = None + + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_called_once_with(op.external_id_key) + task_store.set.assert_called_once_with(op.external_id_key, JOB_NAME) + mock_handle.assert_called_once_with(status="DONE") + + @mock.patch("time.sleep") + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_reconnects_to_running_job_without_resubmitting( + self, mock_hook_cls, mock_submit_job, mock_handle, mock_sleep + ): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = [ + {"status": "RUNNING"}, + {"status": "DONE"}, + ] + + op = _make_operator(poll_interval=5) + task_store = mock.MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = JOB_NAME + + result = op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_called_once_with(op.external_id_key) + task_store.set.assert_not_called() + mock_submit_job.assert_not_called() + mock_handle.assert_called_once_with(status="DONE") + assert result == JOB_NAME + + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "poll_until_complete") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_already_succeeded_completes_without_polling( + self, mock_hook_cls, mock_submit_job, mock_poll_until_complete, mock_handle + ): + mock_hook = mock_hook_cls.return_value + mock_hook.run.side_effect = [{"status": "DONE"}] + + op = _make_operator(poll_interval=5) + task_store = mock.MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = JOB_NAME + + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_called_once_with(op.external_id_key) + mock_submit_job.assert_not_called() + mock_poll_until_complete.assert_not_called() + mock_handle.assert_called_once_with(status="DONE") + + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "poll_until_complete") + @mock.patch.object(SnowparkContainerJobOperator, "get_job_status") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_resubmits_when_stored_job_in_terminal_error( + self, mock_hook_cls, mock_submit_job, mock_get_job_status, mock_poll_until_complete, mock_handle + ): + mock_get_job_status.return_value = "FAILED" + mock_submit_job.return_value = f"{JOB_NAME}_2" + + op = _make_operator(poll_interval=0) + task_store = mock.MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = JOB_NAME + + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_called_once_with(op.external_id_key) + mock_submit_job.assert_called_once() + task_store.set.assert_called_once_with(op.external_id_key, f"{JOB_NAME}_2") + + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "poll_until_complete") + @mock.patch.object(SnowparkContainerJobOperator, "_describe_status") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job") + @mock.patch(MOCK_HOOK_PATH) + def test_resubmits_when_stored_job_not_exist( + self, mock_hook_cls, mock_submit_job, mock_describe_status, mock_poll_until_complete, mock_handle + ): + mock_describe_status.side_effect = ProgrammingError( + msg="test job does not exist", errno=OBJECT_NOT_EXIST_ERROR_CODE + ) + mock_submit_job.return_value = f"{JOB_NAME}_2" + + op = _make_operator(poll_interval=0) + task_store = mock.MagicMock(spec_set=["get", "set"]) + task_store.get.return_value = JOB_NAME + + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_called_once_with(op.external_id_key) + mock_submit_job.assert_called_once() + task_store.set.assert_called_once_with(op.external_id_key, f"{JOB_NAME}_2") + + @mock.patch.object(SnowparkContainerJobOperator, "_handle_final_status") + @mock.patch.object(SnowparkContainerJobOperator, "submit_job", return_value=JOB_NAME) + @mock.patch(MOCK_HOOK_PATH) + def test_durable_false_never_touches_task_state_store(self, mock_hook_cls, mock_submit_job, mock_handle): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = {"status": "DONE"} + + task_store = mock.MagicMock(spec_set=["get", "set"]) + op = _make_operator(durable=False) + + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_not_called() + task_store.set.assert_not_called() + + @mock.patch(MOCK_HOOK_PATH) + def test_deferrable_does_not_persist_to_task_state_store(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = SUBMIT_RESPONSE + + op = _make_operator(deferrable=True) + task_store = mock.MagicMock(spec_set=["get", "set"]) + + with pytest.raises(TaskDeferred): + op.execute(context=_context(task_store=task_store)) + + task_store.get.assert_not_called() + task_store.set.assert_not_called() + + @mock.patch(MOCK_HOOK_PATH) + def test_no_wait_does_not_persist_to_task_state_store(self, mock_hook_cls): + mock_hook = mock_hook_cls.return_value + mock_hook.run.return_value = SUBMIT_RESPONSE + + op = _make_operator(wait_for_completion=False) + task_store = mock.MagicMock(spec_set=["get", "set"]) + + result = op.execute(context=_context(task_store=task_store)) + + assert result == JOB_NAME + task_store.get.assert_not_called() + task_store.set.assert_not_called() + + @mock.patch.object(SnowparkContainerJobOperator, "_describe_status") + def test_get_job_status_returns_status(self, mock_describe_status): + mock_describe_status.return_value = "DONE" + op = _make_operator() + status = op.get_job_status(context=None, external_id="test job") + assert status == "DONE" + + @mock.patch.object(SnowparkContainerJobOperator, "_describe_status") + def test_get_job_status_returns_not_found_when_service_missing(self, mock_describe_status): + mock_describe_status.side_effect = ProgrammingError( + msg="test job does not exist", errno=OBJECT_NOT_EXIST_ERROR_CODE + ) + op = _make_operator() + status = op.get_job_status(context=None, external_id="test job") + assert status == NOT_FOUND_STATUS + + @mock.patch.object(SnowparkContainerJobOperator, "_describe_status") + def test_get_job_status_reraises_other_programming_errors(self, mock_describe_status): + mock_describe_status.side_effect = ProgrammingError(msg="syntax error", errno=1003) + op = _make_operator() + with pytest.raises(ProgrammingError, match="syntax error"): + op.get_job_status(context=None, external_id="test job") + + def test_is_job_active_and_is_job_succeeded(self): + op = _make_operator() + assert op.is_job_active("RUNNING") is True + assert op.is_job_active("DONE") is False + + assert op.is_job_succeeded("DONE") is True + assert op.is_job_succeeded("FAILED") is False + + def test_default_args_durable_reaches_operator(self): + op = _make_operator(default_args={"durable": False}) + assert op.durable is False + + def test_durable_false_direct_kwarg_reaches_operator(self): + op = _make_operator(durable=False) + assert op.durable is False + + +class TestWarnAndDisableDurableAirflowPre3_3: + def test_no_warning_when_unset(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = _warn_and_disable_durable_pre_3_3(_DURABLE_UNSET) + assert result is False + assert caught == [] + + @pytest.mark.parametrize("value", [True, False]) + def test_warns_and_disables_when_explicitly_set(self, value): + with pytest.warns(UserWarning, match="durable.*no effect"): + result = _warn_and_disable_durable_pre_3_3(value) + assert result is False