This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new eb0cf61f7a0 Add deferrable mode to SnowparkContainerJobOperator
(#70103)
eb0cf61f7a0 is described below
commit eb0cf61f7a0e8238abc359f7993449a116c1fab5
Author: Justin Pakzad <[email protected]>
AuthorDate: Tue Sep 8 06:34:55 2026 -0400
Add deferrable mode to SnowparkContainerJobOperator (#70103)
* Add deferrable mode to SnowparkContainerJobOperator
Deferrable mode releases the worker slot while a container job runs, polling
status on the triggerer instead of blocking a worker for the job's full
duration.
The poll is now bounded by a timeout so a job that never reaches a terminal
state no longer blocks indefinitely. The service is dropped on completion,
timeout, or kill to avoid leaving a still-billing resource, and the drop is
best-effort so a cleanup failure does not fail an otherwise successful job.
* Update
providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
---------
Co-authored-by: Jarek Potiuk <[email protected]>
---
providers/snowflake/docs/changelog.rst | 5 +
.../docs/operators/snowpark_containers.rst | 8 +
providers/snowflake/provider.yaml | 1 +
.../providers/snowflake/get_provider_info.py | 5 +-
.../snowflake/operators/snowpark_containers.py | 149 ++++++++++-----
.../snowflake/triggers/snowpark_containers.py | 166 +++++++++++++++++
.../snowflake/utils/snowpark_containers.py | 52 ++++++
.../operators/test_snowpark_containers.py | 153 +++++++++++++++
.../snowflake/triggers/test_snowpark_containers.py | 205 +++++++++++++++++++++
.../snowflake/utils/test_snowpark_containers.py | 28 +++
10 files changed, 723 insertions(+), 49 deletions(-)
diff --git a/providers/snowflake/docs/changelog.rst
b/providers/snowflake/docs/changelog.rst
index 1cf05253c33..938e6395a25 100644
--- a/providers/snowflake/docs/changelog.rst
+++ b/providers/snowflake/docs/changelog.rst
@@ -27,6 +27,11 @@
Changelog
---------
+.. warning::
+ ``SnowparkContainerJobOperator`` now applies a default ``timeout`` of 24
hours where it
+ previously polled indefinitely, so a task running longer than a day now
fails. Increase
+ ``timeout`` to allow more time.
+
6.16.1
......
diff --git a/providers/snowflake/docs/operators/snowpark_containers.rst
b/providers/snowflake/docs/operators/snowpark_containers.rst
index 8677d16207c..3d475350db1 100644
--- a/providers/snowflake/docs/operators/snowpark_containers.rst
+++ b/providers/snowflake/docs/operators/snowpark_containers.rst
@@ -31,6 +31,14 @@ SQL command. It submits the job asynchronously, optionally
polls until the job
reaches a terminal state, retrieves container logs, and can drop the service on
completion.
+You can also run this operator in deferrable mode by setting ``deferrable``
param to ``True``.
+This will ensure that the task is deferred from the Airflow worker slot and
polling for the job status happens on the trigger.
+
+.. note::
+
+ In deferrable mode, clearing or marking a deferred task as failed drops the
job service
+ only on Airflow 3.3 or later. On earlier versions the service is left
running and billing.
+
Prerequisite Tasks
^^^^^^^^^^^^^^^^^^
diff --git a/providers/snowflake/provider.yaml
b/providers/snowflake/provider.yaml
index 415ef25a1b3..31ec68afaf3 100644
--- a/providers/snowflake/provider.yaml
+++ b/providers/snowflake/provider.yaml
@@ -300,6 +300,7 @@ triggers:
- integration-name: Snowflake
python-modules:
- airflow.providers.snowflake.triggers.snowflake_trigger
+ - airflow.providers.snowflake.triggers.snowpark_containers
config:
snowflake:
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
index 1b9456e0298..ea14db2fba0 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
@@ -167,7 +167,10 @@ def get_provider_info():
"triggers": [
{
"integration-name": "Snowflake",
- "python-modules":
["airflow.providers.snowflake.triggers.snowflake_trigger"],
+ "python-modules": [
+ "airflow.providers.snowflake.triggers.snowflake_trigger",
+ "airflow.providers.snowflake.triggers.snowpark_containers",
+ ],
}
],
"config": {
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 ac8e536b920..7ae9175b494 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py
@@ -19,51 +19,25 @@ from __future__ import annotations
import time
from collections.abc import Sequence
-from enum import Enum
+from datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from airflow.providers.common.compat.sdk import conf
from airflow.providers.common.compat.standard.operators import BaseOperator
from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.triggers.snowpark_containers import
SnowparkContainerJobTrigger
+from airflow.providers.snowflake.utils.snowpark_containers import (
+ NON_TERMINAL_STATUSES,
+ TERMINAL_STATUSES,
+ SnowparkContainerJobStatus,
+)
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
-class SnowparkContainerJobStatus(str, Enum):
- """Statuses of a Snowpark Container Services."""
-
- PENDING = "PENDING"
- RUNNING = "RUNNING"
- CANCELLING = "CANCELLING"
- SUSPENDING = "SUSPENDING"
- DELETING = "DELETING"
- DONE = "DONE"
- FAILED = "FAILED"
- CANCELLED = "CANCELLED"
- INTERNAL_ERROR = "INTERNAL_ERROR"
-
-
-TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
- {
- SnowparkContainerJobStatus.DONE,
- SnowparkContainerJobStatus.FAILED,
- SnowparkContainerJobStatus.CANCELLED,
- SnowparkContainerJobStatus.INTERNAL_ERROR,
- }
-)
-NON_TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
- {
- SnowparkContainerJobStatus.PENDING,
- SnowparkContainerJobStatus.RUNNING,
- SnowparkContainerJobStatus.CANCELLING,
- SnowparkContainerJobStatus.SUSPENDING,
- SnowparkContainerJobStatus.DELETING,
- }
-)
-
-
class SnowparkContainerJobOperator(BaseOperator):
"""
Execute a job on Snowpark Container Services.
@@ -94,12 +68,18 @@ class SnowparkContainerJobOperator(BaseOperator):
When disabled, the job is submitted and the operator returns
immediately. (default value: True)
:param drop_on_completion: drop the job service after the job finishes
- successfully. Failed jobs are not dropped, allowing inspection
- in Snowflake. (default value: True)
+ successfully or on a timeout. Failed jobs are not dropped, allowing
+ inspection in Snowflake. (default value: True)
:param poll_interval: the interval in seconds to poll the query status.
(default value: 10)
:param snowflake_conn_id: Reference to
:ref:`Snowflake connection id<howto/connection:snowflake>`
+ :param deferrable: Run the operator in deferrable mode. Only effective when
+ ``wait_for_completion`` is True. With ``wait_for_completion=False`` the
+ operator submits the job and returns immediately without deferring.
+ (default value: False)
+ :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
in connection)
:param schema: name of schema (will overwrite schema defined in
@@ -137,6 +117,8 @@ class SnowparkContainerJobOperator(BaseOperator):
drop_on_completion: bool = True,
poll_interval: int = 10,
snowflake_conn_id: str = "snowflake_default",
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ timeout: int = 24 * 60 * 60,
database: str | None = None,
schema: str | None = None,
role: str | None = None,
@@ -158,6 +140,8 @@ class SnowparkContainerJobOperator(BaseOperator):
self.drop_on_completion = drop_on_completion
self.poll_interval = poll_interval
self.snowflake_conn_id = snowflake_conn_id
+ self.deferrable = deferrable
+ self.timeout = timeout
self.database = database
self.schema = schema
self.role = role
@@ -165,6 +149,9 @@ class SnowparkContainerJobOperator(BaseOperator):
# Set after the job is submitted, parsed from the job submission
response.
self.job_name: str | None = None
+ if self.deferrable and not self.wait_for_completion:
+ self.log.warning("deferrable has no effect when
wait_for_completion is False.")
+
@cached_property
def _hook(self) -> SnowflakeHook:
return SnowflakeHook(
@@ -203,7 +190,14 @@ class SnowparkContainerJobOperator(BaseOperator):
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ status = None
+ end_time = time.monotonic() + self.timeout
while True:
+ if time.monotonic() >= end_time:
+ self._log_container_output(status)
+ 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")
if status in TERMINAL_STATUSES:
@@ -212,11 +206,15 @@ class SnowparkContainerJobOperator(BaseOperator):
raise RuntimeError(f"Job {self.job_name} returned unexpected
status: {status}")
time.sleep(self.poll_interval)
- def _log_container_output(self, status: str) -> None:
- """Fetch and log container output for all replicas."""
+ 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):
sql = f"SELECT SYSTEM$GET_SERVICE_LOGS('{self.job_name}',
{instance_id}, '{self.container_name}')"
- response = self._run_one(sql)[0]
+ try:
+ response = self._run_one(sql)[0]
+ except Exception as e:
+ self.log.warning("Could not retrieve logs for instance_id %d:
%s", instance_id, e)
+ continue
if not response:
continue
if status != SnowparkContainerJobStatus.DONE:
@@ -224,13 +222,27 @@ class SnowparkContainerJobOperator(BaseOperator):
else:
self.log.info("Logs for instance_id %d:\n%s", instance_id,
response)
+ def _drop_service(self) -> None:
+ """Best-effort drop of the job service."""
+ try:
+ self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
+ except Exception as e:
+ self.log.error("Error dropping service %s: %s", self.job_name, e)
+
def on_kill(self) -> None:
"""Drop the running service on task kill."""
if self.job_name:
- try:
- self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
- except Exception as e:
- self.log.error("Error dropping service %s: %s", self.job_name,
e)
+ self._drop_service()
+
+ def _handle_final_status(self, status: str) -> None:
+ """Log container output, fail unless the job is DONE, and optionally
drop the service on success."""
+ self._log_container_output(status)
+ if status != SnowparkContainerJobStatus.DONE:
+ raise RuntimeError(f"Job '{self.job_name}' finished with status:
{status}")
+ if self.drop_on_completion:
+ # Job already succeeded, so a cleanup failure is logged rather
than raised
+ # to avoid marking a successful job as failed.
+ self._drop_service()
def execute(self, context: Context) -> str:
"""Submit and optionally wait for a Snowpark Container Services job."""
@@ -242,10 +254,51 @@ class SnowparkContainerJobOperator(BaseOperator):
raise RuntimeError("Job name was not returned")
if not self.wait_for_completion:
return self.job_name
+ if self.deferrable:
+ # 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.
+ now = time.time()
+ poll_buffer = timedelta(seconds=self.poll_interval + 60)
+ execution_deadline = None
+ defer_timeout = timedelta(seconds=self.timeout) + poll_buffer
+ if self.execution_timeout is not None:
+ # Hand the execution deadline to the trigger so it emits a
timeout event that drops the
+ # service. The framework's defer timeout would otherwise kill
the task with no cleanup.
+ execution_deadline = (
+ context["ti"].start_date.timestamp() +
self.execution_timeout.total_seconds()
+ )
+ # Pad the backstop past that deadline so the trigger fires
first.
+ defer_timeout = self.execution_timeout + poll_buffer
+ self.defer(
+ trigger=SnowparkContainerJobTrigger(
+ job_name=self.job_name,
+ snowflake_conn_id=self.snowflake_conn_id,
+ poll_interval=self.poll_interval,
+ end_time=now + self.timeout,
+ execution_deadline=execution_deadline,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ warehouse=self.warehouse,
+ ),
+ timeout=defer_timeout,
+ method_name="execute_complete",
+ )
status = self._poll_for_status()
- self._log_container_output(status)
- if status != SnowparkContainerJobStatus.DONE:
- raise RuntimeError(f"Job '{self.job_name}' finished with status:
{status}")
- if self.drop_on_completion:
- self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
+ self._handle_final_status(status)
+ return self.job_name
+
+ def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
+ """Resume after the trigger fires."""
+ self.job_name = event["job_name"]
+ status = event["status"]
+ if status == "timeout":
+ self._log_container_output(status)
+ if self.drop_on_completion:
+ self._drop_service()
+ raise TimeoutError(event.get("message", f"Job '{self.job_name}'
did not complete: {status}"))
+ if status == "error":
+ raise RuntimeError(event.get("message", f"Job '{self.job_name}'
did not complete: {status}"))
+ self._handle_final_status(status)
return self.job_name
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py
new file mode 100644
index 00000000000..4b1c6e2d1be
--- /dev/null
+++
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py
@@ -0,0 +1,166 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.utils.snowpark_containers import (
+ NON_TERMINAL_STATUSES,
+ TERMINAL_STATUSES,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class SnowparkContainerJobTrigger(BaseTrigger):
+ """
+ Poll a Snowpark Container Services job until it reaches a terminal status.
+
+ :param job_name: name of the submitted job service to poll.
+ :param snowflake_conn_id: reference to the Snowflake connection id.
+ :param poll_interval: seconds to sleep between ``DESCRIBE SERVICE`` polls.
+ :param end_time: epoch deadline (``time.time()`` seconds) after which a
``timeout``
+ event is emitted.
+ :param execution_deadline: (Optional) absolute timestamp (in seconds since
the epoch) after
+ which the task is considered timed out. (default: None)
+ :param database: (Optional) name of database. (default: None)
+ :param schema: (Optional) name of schema. (default: None)
+ :param role: (Optional) name of role. (default: None)
+ :param warehouse: (Optional) name of warehouse. (default: None)
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ execution_deadline: float | None = None,
+ database: str | None = None,
+ schema: str | None = None,
+ role: str | None = None,
+ warehouse: str | None = None,
+ ) -> None:
+ super().__init__()
+ self.job_name = job_name
+ self.snowflake_conn_id = snowflake_conn_id
+ self.poll_interval = poll_interval
+ self.end_time = end_time
+ self.execution_deadline = execution_deadline
+ self.database = database
+ self.schema = schema
+ self.role = role
+ self.warehouse = warehouse
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize SnowparkContainerJobTrigger arguments and class path."""
+ return (
+
"airflow.providers.snowflake.triggers.snowpark_containers.SnowparkContainerJobTrigger",
+ {
+ "job_name": self.job_name,
+ "snowflake_conn_id": self.snowflake_conn_id,
+ "poll_interval": self.poll_interval,
+ "end_time": self.end_time,
+ "execution_deadline": self.execution_deadline,
+ "database": self.database,
+ "schema": self.schema,
+ "role": self.role,
+ "warehouse": self.warehouse,
+ },
+ )
+
+ def _get_hook(self) -> SnowflakeHook:
+ """Build a ``SnowflakeHook`` from the trigger's connection settings."""
+ return SnowflakeHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ warehouse=self.warehouse,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ )
+
+ async def _describe_status(self, hook: SnowflakeHook) -> str | None:
+ """Return the job's current status via ``DESCRIBE SERVICE``, or
``None`` if absent."""
+ # SnowflakeHook is synchronous. Run the blocking poll off the event
loop so a
+ # single query does not stall every other trigger on this triggerer.
+ response: Any = await asyncio.to_thread(
+ hook.run,
+ f"DESCRIBE SERVICE {self.job_name}",
+ handler=fetch_one_handler,
+ return_dictionaries=True,
+ )
+ return response.get("status") if response else None
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Poll the job status and yield exactly one terminal event."""
+ hook = self._get_hook()
+ while True:
+ now = time.time()
+ if self.execution_deadline is not None and now >=
self.execution_deadline:
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} reached the execution
timeout.",
+ }
+ )
+ return
+
+ if now >= self.end_time:
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} did not reach a
terminal status before the timeout.",
+ }
+ )
+ return
+
+ try:
+ status = await self._describe_status(hook=hook)
+ except Exception as e:
+ yield TriggerEvent({"status": "error", "job_name":
self.job_name, "message": str(e)})
+ return
+
+ if status in TERMINAL_STATUSES:
+ yield TriggerEvent({"status": status, "job_name":
self.job_name})
+ return
+
+ if status not in NON_TERMINAL_STATUSES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} returned unexpected
status: {status}",
+ }
+ )
+ return
+
+ await asyncio.sleep(self.poll_interval)
+
+ async def on_kill(self) -> None:
+ """Drop the job service when a deferred task is killed."""
+ hook = self._get_hook()
+ try:
+ await asyncio.to_thread(hook.run, f"DROP SERVICE IF EXISTS
{self.job_name}")
+ self.log.info("on_kill: dropped service %s", self.job_name)
+ except Exception as e:
+ self.log.error("on_kill: failed to drop service %s: %s",
self.job_name, e)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py
b/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py
new file mode 100644
index 00000000000..001446123b1
--- /dev/null
+++
b/providers/snowflake/src/airflow/providers/snowflake/utils/snowpark_containers.py
@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from enum import Enum
+
+
+class SnowparkContainerJobStatus(str, Enum):
+ """Statuses of a Snowpark Container Services job service."""
+
+ PENDING = "PENDING"
+ RUNNING = "RUNNING"
+ CANCELLING = "CANCELLING"
+ SUSPENDING = "SUSPENDING"
+ DELETING = "DELETING"
+ DONE = "DONE"
+ FAILED = "FAILED"
+ CANCELLED = "CANCELLED"
+ INTERNAL_ERROR = "INTERNAL_ERROR"
+
+
+TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
+ {
+ SnowparkContainerJobStatus.DONE,
+ SnowparkContainerJobStatus.FAILED,
+ SnowparkContainerJobStatus.CANCELLED,
+ SnowparkContainerJobStatus.INTERNAL_ERROR,
+ }
+)
+NON_TERMINAL_STATUSES: frozenset[SnowparkContainerJobStatus] = frozenset(
+ {
+ SnowparkContainerJobStatus.PENDING,
+ SnowparkContainerJobStatus.RUNNING,
+ SnowparkContainerJobStatus.CANCELLING,
+ SnowparkContainerJobStatus.SUSPENDING,
+ SnowparkContainerJobStatus.DELETING,
+ }
+)
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 c154c529025..12262c22901 100644
---
a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
+++
b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py
@@ -16,11 +16,14 @@
# under the License.
from __future__ import annotations
+from datetime import datetime, timedelta, timezone
from unittest import mock
import pytest
+from airflow.providers.common.compat.sdk import TaskDeferred
from airflow.providers.snowflake.operators.snowpark_containers import
SnowparkContainerJobOperator
+from airflow.providers.snowflake.triggers.snowpark_containers import
SnowparkContainerJobTrigger
TASK_ID = "test_spcs_job"
COMPUTE_POOL = "test_pool"
@@ -75,6 +78,21 @@ class TestSnowparkContainerJobOperator:
with pytest.raises(ValueError, match=match):
op.execute(context=None)
+ @pytest.mark.parametrize(
+ ("deferrable", "wait_for_completion", "warns"),
+ (
+ pytest.param(True, False, True, id="deferrable_no_wait"),
+ pytest.param(True, True, False, id="deferrable_wait"),
+ pytest.param(False, False, False, id="sync_no_wait"),
+ ),
+ )
+ @mock.patch.object(SnowparkContainerJobOperator, "log")
+ def test_warns_when_deferrable_without_wait_for_completion(
+ self, mock_log, deferrable, wait_for_completion, warns
+ ):
+ _make_operator(deferrable=deferrable,
wait_for_completion=wait_for_completion)
+ assert mock_log.warning.called is warns
+
def test_build_sql_with_spec_stage(self):
op = _make_operator()
sql = op._build_sql()
@@ -150,6 +168,32 @@ class TestSnowparkContainerJobOperator:
assert op._poll_for_status() == "DONE"
assert mock_sleep.call_count == 2
+ @pytest.mark.parametrize(
+ ("drop_on_completion", "drops"),
+ [(True, True), (False, False)],
+ )
+ @mock.patch("time.sleep")
+ @mock.patch("time.monotonic", side_effect=itertools.count(0, 20))
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_poll_raises_and_logs_on_timeout(
+ self, mock_log, mock_hook_cls, mock_monotonic, sleep_mock,
drop_on_completion, drops
+ ):
+ mock_hook = mock_hook_cls.return_value
+ mock_hook.run.return_value = {"status": "RUNNING"}
+ op = _make_operator(poll_interval=5, timeout=10,
drop_on_completion=drop_on_completion)
+ op.job_name = JOB_NAME
+
+ with pytest.raises(TimeoutError, match="did not reach a terminal
status"):
+ op._poll_for_status()
+
+ mock_log.assert_called_once_with("RUNNING")
+ drop_call = mock.call(f"DROP SERVICE IF EXISTS {JOB_NAME}")
+ if drops:
+ assert drop_call in mock_hook.run.call_args_list
+ else:
+ assert drop_call not in mock_hook.run.call_args_list
+
@mock.patch(MOCK_HOOK_PATH)
def test_log_container_output_uses_info_on_done(self, mock_hook_cls):
mock_hook = mock_hook_cls.return_value
@@ -191,6 +235,16 @@ class TestSnowparkContainerJobOperator:
op._log_container_output("DONE")
assert mock_info.call_count == 3
+ @mock.patch(MOCK_HOOK_PATH)
+ def test_log_container_output_swallows_fetch_error(self, mock_hook_cls):
+ mock_hook = mock_hook_cls.return_value
+ mock_hook.run.side_effect = Exception("Unable to retrieve logs")
+ op = _make_operator()
+ op.job_name = JOB_NAME
+ with mock.patch.object(op.log, "warning") as mock_warning:
+ op._log_container_output("RUNNING")
+ mock_warning.assert_called_once_with("Could not retrieve logs for
instance_id %d: %s", 0, mock.ANY)
+
@mock.patch(MOCK_HOOK_PATH)
def test_on_kill_no_job_name(self, mock_hook_cls):
mock_hook = mock_hook_cls.return_value
@@ -280,3 +334,102 @@ class TestSnowparkContainerJobOperator:
op = _make_operator(drop_on_completion=False)
op.execute(context=None)
mock_hook.run.assert_not_called()
+
+ @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:
+ op.execute(context=None)
+ assert isinstance(exc.value.trigger, SnowparkContainerJobTrigger)
+ 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)
+ 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:
+ op.execute(context=None)
+ 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)
+ 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))}
+ op = _make_operator(
+ deferrable=True,
+ timeout=3600,
+ poll_interval=10,
+ execution_timeout=timedelta(seconds=120),
+ )
+ with pytest.raises(TaskDeferred) as exc:
+ op.execute(context=context)
+ assert exc.value.trigger.end_time == 1000 + 3600
+ assert exc.value.trigger.execution_deadline == 1000 + 120
+ assert exc.value.timeout == timedelta(seconds=120 + 10 + 60)
+
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_execute_complete_success_drops_and_returns(self, mock_log,
mock_hook_cls):
+ mock_hook = mock_hook_cls.return_value
+ op = _make_operator(drop_on_completion=True)
+ result = op.execute_complete(context=None, event={"status": "DONE",
"job_name": JOB_NAME})
+ assert result == JOB_NAME
+ mock_log.assert_called_once_with("DONE")
+ mock_hook.run.assert_called_once_with(f"DROP SERVICE IF EXISTS
{JOB_NAME}")
+
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_execute_complete_failure_raises_without_drop(self, mock_log,
mock_hook_cls):
+ mock_hook = mock_hook_cls.return_value
+ op = _make_operator()
+ with pytest.raises(RuntimeError, match="FAILED"):
+ op.execute_complete(context=None, event={"status": "FAILED",
"job_name": JOB_NAME})
+ mock_log.assert_called_once_with("FAILED")
+ mock_hook.run.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("status", "exc", "drops", "logs"),
+ [("timeout", TimeoutError, True, True), ("error", RuntimeError, False,
False)],
+ )
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_execute_complete_raises_and_drops_only_on_timeout(
+ self, mock_log, mock_hook_cls, status, exc, drops, logs
+ ):
+ mock_hook = mock_hook_cls.return_value
+ op = _make_operator()
+ with pytest.raises(exc, match="boom"):
+ op.execute_complete(
+ context=None,
+ event={"status": status, "job_name": JOB_NAME, "message":
"boom"},
+ )
+ if drops:
+ mock_hook.run.assert_called_once_with(f"DROP SERVICE IF EXISTS
{JOB_NAME}")
+ else:
+ mock_hook.run.assert_not_called()
+ if logs:
+ mock_log.assert_called_once_with(status)
+ else:
+ mock_log.assert_not_called()
+
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_execute_complete_skips_drop_when_disabled(self, mock_log,
mock_hook_cls):
+ mock_hook = mock_hook_cls.return_value
+ op = _make_operator(drop_on_completion=False)
+ op.execute_complete(context=None, event={"status": "DONE", "job_name":
JOB_NAME})
+ mock_hook.run.assert_not_called()
+
+ @mock.patch(MOCK_HOOK_PATH)
+ @mock.patch.object(SnowparkContainerJobOperator, "_log_container_output")
+ def test_execute_complete_timeout_skips_drop_when_disabled(self, mock_log,
mock_hook_cls):
+ mock_hook = mock_hook_cls.return_value
+ op = _make_operator(drop_on_completion=False)
+ with pytest.raises(TimeoutError, match="boom"):
+ op.execute_complete(
+ context=None,
+ event={"status": "timeout", "job_name": JOB_NAME, "message":
"boom"},
+ )
+ mock_log.assert_called_once_with("timeout")
+ mock_hook.run.assert_not_called()
diff --git
a/providers/snowflake/tests/unit/snowflake/triggers/test_snowpark_containers.py
b/providers/snowflake/tests/unit/snowflake/triggers/test_snowpark_containers.py
new file mode 100644
index 00000000000..f0e29e4515d
--- /dev/null
+++
b/providers/snowflake/tests/unit/snowflake/triggers/test_snowpark_containers.py
@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import time
+from unittest import mock
+
+import pytest
+
+from airflow.providers.snowflake.triggers.snowpark_containers import
SnowparkContainerJobTrigger
+from airflow.triggers.base import TriggerEvent
+
+TRIGGER_PATH = "airflow.providers.snowflake.triggers.snowpark_containers"
+CLASSPATH = f"{TRIGGER_PATH}.SnowparkContainerJobTrigger"
+HOOK = f"{TRIGGER_PATH}.SnowflakeHook"
+
+JOB_NAME = "TEST_JOB"
+CONN_ID = "snowflake_default"
+POLL_INTERVAL = 1.0
+
+
+class TestSnowparkContainerJobTrigger:
+ @staticmethod
+ def _describe(status):
+ return {"status": status}
+
+ def _trigger(self, end_time=None, **kwargs):
+ params = {
+ "job_name": JOB_NAME,
+ "snowflake_conn_id": CONN_ID,
+ "poll_interval": POLL_INTERVAL,
+ "end_time": end_time if end_time is not None else time.time() +
3600,
+ }
+ params.update(kwargs)
+ return SnowparkContainerJobTrigger(**params)
+
+ def test_serialization(self):
+ end_time = time.time() + 3600
+ execution_deadline = time.time() + 1800
+ class_path, kwargs = self._trigger(
+ end_time,
+ execution_deadline=execution_deadline,
+ database="db",
+ schema="sc",
+ role="r",
+ warehouse="wh",
+ ).serialize()
+ assert class_path == CLASSPATH
+ assert kwargs == {
+ "job_name": JOB_NAME,
+ "snowflake_conn_id": CONN_ID,
+ "poll_interval": POLL_INTERVAL,
+ "end_time": end_time,
+ "execution_deadline": execution_deadline,
+ "database": "db",
+ "schema": "sc",
+ "role": "r",
+ "warehouse": "wh",
+ }
+
+ @mock.patch(HOOK, autospec=True)
+ def test_get_hook_builds_hook_from_connection_settings(self,
mock_hook_cls):
+ self._trigger(database="db", schema="sc", role="r",
warehouse="wh")._get_hook()
+ mock_hook_cls.assert_called_once_with(
+ snowflake_conn_id=CONN_ID,
+ warehouse="wh",
+ database="db",
+ schema="sc",
+ role="r",
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_describe_status_returns_none_when_no_row(self,
mock_hook_cls):
+ mock_hook_cls.return_value.run.return_value = None
+ trigger = self._trigger()
+ assert await trigger._describe_status(trigger._get_hook()) is None
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("status", ("DONE", "FAILED", "CANCELLED",
"INTERNAL_ERROR"))
+ @mock.patch(HOOK, autospec=True)
+ async def test_terminal_status_yields_status_event(self, mock_hook_cls,
status):
+ mock_hook_cls.return_value.run.return_value = self._describe(status)
+ event = await self._trigger().run().__anext__()
+ assert event == TriggerEvent({"status": status, "job_name": JOB_NAME})
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_unexpected_status_yields_error(self, mock_hook_cls):
+ mock_hook_cls.return_value.run.return_value = self._describe("RANDOM")
+ event = await self._trigger().run().__anext__()
+ assert event.payload["status"] == "error"
+ assert "unexpected status" in event.payload["message"]
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_timeout_yields_timeout_event(self, mock_hook_cls):
+ mock_hook_cls.return_value.run.return_value = self._describe("RUNNING")
+ event = await self._trigger(end_time=time.time() - 1).run().__anext__()
+ assert event.payload["status"] == "timeout"
+ assert event.payload["job_name"] == JOB_NAME
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_execution_deadline_yields_timeout(self, mock_hook_cls):
+ mock_hook_cls.return_value.run.return_value = self._describe("RUNNING")
+ trigger = self._trigger(end_time=time.time() + 3600,
execution_deadline=time.time() - 1)
+ event = await trigger.run().__anext__()
+ assert event.payload["status"] == "timeout"
+ assert "execution timeout" in event.payload["message"]
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("end_time", "execution_deadline"),
+ (
+ (500.0, None),
+ (2000.0, 500.0),
+ ),
+ )
+ @mock.patch(HOOK, autospec=True)
+ async def test_deadline_takes_precedence_over_terminal_status(
+ self, mock_hook_cls, end_time, execution_deadline, time_machine
+ ):
+ time_machine.move_to(1000, tick=False)
+ mock_hook_cls.return_value.run.return_value = self._describe("DONE")
+ trigger = self._trigger(end_time=end_time,
execution_deadline=execution_deadline)
+ event = await trigger.run().__anext__()
+ assert event.payload["status"] == "timeout"
+ mock_hook_cls.return_value.run.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("end_time", "execution_deadline"),
+ (
+ (1000.0, None),
+ (2000.0, 1000.0),
+ ),
+ )
+ @mock.patch(HOOK, autospec=True)
+ async def test_deadline_boundary_counts_as_expired(
+ self, mock_hook_cls, end_time, execution_deadline, time_machine
+ ):
+ time_machine.move_to(1000, tick=False)
+ mock_hook_cls.return_value.run.return_value = self._describe("RUNNING")
+ trigger = self._trigger(end_time=end_time,
execution_deadline=execution_deadline)
+ event = await trigger.run().__anext__()
+ assert event.payload["status"] == "timeout"
+ mock_hook_cls.return_value.run.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_future_execution_deadline_does_not_short_circuit(self,
mock_hook_cls):
+ mock_hook_cls.return_value.run.return_value = self._describe("DONE")
+ trigger = self._trigger(end_time=time.time() + 3600,
execution_deadline=time.time() + 1800)
+ event = await trigger.run().__anext__()
+ assert event == TriggerEvent({"status": "DONE", "job_name": JOB_NAME})
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_poll_exception_yields_error(self, mock_hook_cls):
+ mock_hook_cls.return_value.run.side_effect = RuntimeError("boom")
+ event = await self._trigger().run().__anext__()
+ assert event == TriggerEvent({"status": "error", "job_name": JOB_NAME,
"message": "boom"})
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep")
+ @mock.patch(HOOK, autospec=True)
+ async def test_polls_through_non_terminal_then_terminal(self,
mock_hook_cls, mock_sleep):
+ mock_hook_cls.return_value.run.side_effect = [
+ self._describe("PENDING"),
+ self._describe("RUNNING"),
+ self._describe("DONE"),
+ ]
+ event = await self._trigger().run().__anext__()
+ assert event == TriggerEvent({"status": "DONE", "job_name": JOB_NAME})
+ assert mock_sleep.await_count == 2
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_on_kill_drops_service(self, mock_hook_cls):
+ await self._trigger().on_kill()
+ mock_hook_cls.return_value.run.assert_called_once_with(f"DROP SERVICE
IF EXISTS {JOB_NAME}")
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK, autospec=True)
+ async def test_on_kill_logs_error_on_failure(self, mock_hook_cls):
+ mock_hook_cls.return_value.run.side_effect = RuntimeError("drop
failed")
+ trigger = self._trigger()
+ with mock.patch.object(trigger.log, "error") as mock_error:
+ await trigger.on_kill()
+ mock_error.assert_called_once()
diff --git
a/providers/snowflake/tests/unit/snowflake/utils/test_snowpark_containers.py
b/providers/snowflake/tests/unit/snowflake/utils/test_snowpark_containers.py
new file mode 100644
index 00000000000..3358d9ea782
--- /dev/null
+++ b/providers/snowflake/tests/unit/snowflake/utils/test_snowpark_containers.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow.providers.snowflake.utils.snowpark_containers import (
+ NON_TERMINAL_STATUSES,
+ TERMINAL_STATUSES,
+ SnowparkContainerJobStatus,
+)
+
+
+def test_each_status_is_terminal_or_non_terminal():
+ assert TERMINAL_STATUSES.isdisjoint(NON_TERMINAL_STATUSES)
+ assert set(SnowparkContainerJobStatus) == TERMINAL_STATUSES |
NON_TERMINAL_STATUSES
