You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "justinpakzad (via GitHub)" <gi...@apache.org> on 2026/07/19 20:47:07 UTC
[PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
justinpakzad opened a new pull request, #70103:
URL: https://github.com/apache/airflow/pull/70103
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
**Description**
Follow-up to #68259 adding deferrable mode to SnowparkContainerJobOperator. Snowpark Container jobs are often long-running (ML training, batch processing, other containerized workloads that run for hours), which makes them a good fit for deferrable mode.
The SnowflakeHook is synchronous, so this follows a similar approach the Anthropic provider took. The trigger wraps each blocking hook call in asyncio.to_thread so a single poll doesn't stall the shared triggerer event loop. A timeout parameter is introduced for both modes. Cleanup is consistent across both modes, on timeout the service is dropped to stop billing on a job past its deadline. On poll errors and terminal failures it's left in place for inspection.
The status Enum and terminal/non-terminal frozensets were moved to the trigger file so the operator can import them without a circular dependency.
**Testing**
- Unit tests for the trigger (serialization, terminal/unexpected/timeout/error events, cleanup) and the operator's defer and execute_complete paths.
- Manually tested end-to-end against a real Snowflake account: deferred success (job completes, logs retrieved, service dropped), terminal failure, both timeout types, on_kill, and the drop toggles.
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [x] Yes (Claude Opus 4.6)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890928483
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
While I was working on a follow up to implement retries (so a transient failure doesn't terminate the trigger/operator), I realized the SPCS trigger and operator will need some shared helpers beyond the statuses. So instead of putting them in the hook, I've moved them into a util new file in the existing `/util` directory. This avoids the need to alias the status imports to maintain backwards compatibility and enables adding any helpers that are only used for SPCS.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,12 +191,16 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ end_time = time.time() + self.timeout
while True:
+ if time.time() >= end_time:
+ self._drop_service()
Review Comment:
Good call. Both timeout branches now log the container output. `_log_container_output` is now best effort (wrapped in try/except), since any failures fetching the logs would prevent the cleanup. As for your second point, I've added the guard to the timeout branches as well so the `drop_on_completion` is respected for timeouts. Tests have been added and updated accordingly.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -19,51 +19,25 @@
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.hooks.snowflake import (
+ CONTAINER_JOB_NON_TERMINAL_STATUSES,
Review Comment:
See my comment below. I've moved the statuses to a new utility file and reverted back to the original naming since they are not ambiguous anymore.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3744929476
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
Review Comment:
These statuses describe the Snowpark job/service rather than the trigger itself, and they're also consumed by the operator. Would it make more sense to keep the enum/status sets with the hook (similar to dbt Cloud), so the operator module does not need to import shared abstractions from the trigger module?
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
Review Comment:
I think you should add the default value for each parameter to the docstring entries.
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py:
##########
@@ -280,3 +312,58 @@ def test_execute_skips_drop_when_disabled(self, mock_hook_cls, mock_submit, mock
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"
Review Comment:
If you implement the deadline handling suggested above, could we add coverage for it here? In particular, an explicit `execution_timeout` should determine the trigger's execution deadline, while the timeout passed to `defer()` should include the additional polling buffer.
##########
providers/snowflake/tests/unit/snowflake/triggers/test_snowpark_containers.py:
##########
@@ -0,0 +1,143 @@
+# 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:
Review Comment:
If you implement the suggested deadline handling in the trigger above, these tests will need updating as well to verify that the timeout takes precedence once the deadline has been reached, including over a terminal status.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -242,10 +247,35 @@ def execute(self, context: Context) -> str:
raise RuntimeError("Job name was not returned")
if not self.wait_for_completion:
return self.job_name
+ if self.deferrable:
+ self.defer(
+ trigger=SnowparkContainerJobTrigger(
+ job_name=self.job_name,
+ snowflake_conn_id=self.snowflake_conn_id,
+ poll_interval=self.poll_interval,
+ end_time=time.time() + self.timeout,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ warehouse=self.warehouse,
+ ),
+ # Pad past the trigger's end_time so its timeout event, which drops the service,
+ # fires before this hard backstop. A user-set execution_timeout takes precedence.
+ timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60),
Review Comment:
I don't think the timeout passed to `defer() `is necessarily consistent with the deadline used by the trigger. The trigger's `end_time` is always derived from self.timeout, whereas `defer(timeout=...) `uses self.execution_timeout when one is provided. For example, `execution_timeout=60s` with the `default timeout=86400s` gives the framework a 60s timeout while the trigger still has a 24h deadline. This may result in an orphaned job as the trigger will timeout before the timeout event can be emitted (which will be caught by `execute_complete` to trigger service cancellation). I know you added an on_kill method to potentially handle this situation but looking at the BaseTrigger class, it appears that this specific method is inert in timeout scenarios.
We could pass the execution deadline to the trigger as well and let it determine the applicable deadline, (please see what has been done for the Airbyte provider). The timeout passed to `defer()` can then be padded slightly so the trigger has an opportunity to emit the corresponding timeout event first.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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
Review Comment:
Just checking if 'status' is guaranteed to be upper case in the response object.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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:
+ try:
+ status = await self._describe_status(hook=hook)
+ except Exception as e:
Review Comment:
Do we want a single transient `DESCRIBE SERVICE `failure to terminate the trigger, given these jobs may run for hours, or should polling failures be retried up to the deadline?
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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:
+ 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
+
+ if time.time() > self.end_time:
Review Comment:
1) Change this to `>=` to handle timeouts on the boundary.
2) Should the timeout check happen before polling/evaluating the job status? Currently, if the trigger wakes after `end_time` (or the `DESCRIBE SERVICE` call itself takes us past it) and the service returns a terminal status, we'll emit that terminal event even though the trigger deadline has already expired. If `end_time` is intended to be a hard deadline, I think it should be checked before starting another poll.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5288255966
> Just one nit. Please check my response to the enum comment too.
Addressed both of the comments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3754551726
##########
providers/snowflake/tests/unit/snowflake/triggers/test_snowpark_containers.py:
##########
@@ -0,0 +1,143 @@
+# 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:
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3956945221
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py:
##########
@@ -150,6 +168,32 @@ def test_poll_waits_through_pending_then_done(self, mock_hook_cls, mock_sleep):
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=[0, 0, 20])
Review Comment:
`side_effect=[0, 0, 20]` supplies exactly the three values `_poll_for_status` consumes today: one for `end_time`, one for the check that must be under the deadline, one for the check that must be over it. It works, and the parametrized drop/no-drop split still exercises both cleanup paths.
The fragility is that `mock.patch("time.monotonic", ...)` patches the clock **process-wide** for the whole test body, not just the call under test. Anything else that reads the monotonic clock while the patch is active — `_make_operator`'s `BaseOperator.__init__`, a logging handler, or simply a fourth call if the loop ever grows one — raises `StopIteration` from somewhere with no obvious connection to the assertion.
An unbounded counter reads the same and cannot exhaust:
```suggestion
@mock.patch("time.monotonic", side_effect=itertools.count(0, 20))
```
(with `import itertools` at the top). Same three values consumed, same behaviour, no coupling to the exact call count.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890286076
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
> So I had originally placed them in the trigger as you mentioned, but it was pointed it out in one of the review comments above that it would make more sense to keep the enum/status sets with the hook as importing shared abstractions from the trigger module might look a bit strange. I think it's fine where it is but I'm happy to move em back to the trigger if you feel they are suited better there.
I think it is ultimately down to individual preferences. I do not feel strongly about the placement even though I made that comment. It’s up to you really
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #70103:
URL: https://github.com/apache/airflow/pull/70103
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890135508
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
So I had originally placed them in the trigger as you mentioned, but it was pointed it out in one of the review comments above that it would make more sense to keep the enum/status sets with the hook as importing shared abstractions from the trigger module might look a bit strange. I think it's fine where it is but I'm happy to move em back to the trigger if you feel they are suited better there.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5576435090
> LGTM — all four points from my last pass are addressed: the deadline now derives from `ti.start_date`, `TERMINAL_STATUSES` / `NON_TERMINAL_STATUSES` keep their original names and stay importable from the operator module, both timeout paths log container output before dropping and respect `drop_on_completion`, and the statuses moved out of the SQL hook module into `utils/snowpark_containers.py`. CI is green and every thread is resolved. Four smaller observations inline, none blocking.
Appreciate the second round of review. All good callouts and they've been addressed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3779009509
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -248,20 +247,29 @@ def execute(self, context: Context) -> str:
if not self.wait_for_completion:
return self.job_name
if self.deferrable:
Review Comment:
I would add a comment here explaining the interaction of `execution_timeout` and `timeout` i.e. that the former takes precedence over the latter if both are present.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890928483
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
While I was working on a follow up to implement retries (so a transient failure doesn't terminate the trigger/operator), I realized the SPCS trigger and operator will need some shared helpers beyond the statuses. So instead of putting them in the hook, I've moved them into a new util file in the existing `/utils` directory. This avoids the need to alias the status imports to maintain backwards compatibility and enables adding new helpers that are only used for SPCS.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890929736
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -242,10 +246,47 @@ def execute(self, context: Context) -> str:
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 = now + self.execution_timeout.total_seconds()
Review Comment:
Good find. Updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3952908139
##########
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:
Review Comment:
Good point. Added a note in the `snowpark_containers.rst`.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -212,25 +206,41 @@ def _poll_for_status(self) -> str:
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:
self.log.error("Logs for instance_id %d:\n%s", instance_id, response)
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:
+ self._drop_service()
Review Comment:
Good call out. The switch to best-effort was intentional and I'm aware it does change the current behavior but I think it's the better approach here. A job that already succeeded could be marked as failed (due to a transient error while dropping) and end up retrying an already successful job. I put a comment explaining why.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3952257954
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,7 +190,14 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ status = None
+ end_time = time.time() + self.timeout
Review Comment:
> - `time.monotonic()` for durations, not `time.time()`.
(`AGENTS.md` → Coding Standards)
This loop is entirely in-process, so `time.monotonic()` is the standard-compliant choice and is immune to a wall-clock step (NTP correction, manual clock change) silently shortening or extending the window. It'd be `time.monotonic()` on both this line and the `time.time() >= end_time` check below.
The trigger's `end_time` / `execution_deadline` are a different case and correctly stay on `time.time()` — those are absolute deadlines that get serialized and re-read in another process, where a monotonic value would be meaningless.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -137,6 +117,8 @@ def __init__(
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,
Review Comment:
This default also changes the existing non-deferrable path: `_poll_for_status` polled indefinitely in 6.15.0 / 6.16.1, and now a job still running at 24h fails the task and — with the default `drop_on_completion=True` — drops the service out from under it.
That's a sensible default, but it's a behaviour change for anyone already running SPCS jobs longer than a day, and it isn't visible from the PR title. Airflow's provider release managers regenerate `changelog.rst` from `git log`, so a line in `providers/snowflake/docs/changelog.rst` (or at minimum a sentence in the commit body) is the only place it'll get flagged for them.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -212,25 +206,41 @@ def _poll_for_status(self) -> str:
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:
self.log.error("Logs for instance_id %d:\n%s", instance_id, response)
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:
+ self._drop_service()
Review Comment:
Small behaviour change worth being deliberate about. Before this PR the success-path drop was:
```python
if self.drop_on_completion:
self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
```
— no `try`, so a drop that failed failed the task. Routing it through `_drop_service` (line 225) now logs at `error` level and lets the task succeed, which leaves a service running and billing with no failed task to alert on.
Swallowing may well be what you want here — the job did succeed, and failing an otherwise-good task on a cleanup error is harsh. But `on_kill` was best-effort because it runs during teardown, whereas this is the normal completion path, and the failure is now silent. Either keeping the raise on the completion path or a one-line note on why it's best-effort would make the choice legible.
##########
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:
Review Comment:
`BaseTrigger.on_kill()` only landed in Airflow 3.3.0 (apache/airflow#65590), and this provider's floor is `apache-airflow>=2.11.0`. On 2.11 through 3.2 nothing calls this method, so clearing or marking-failed a deferred SPCS task leaves the job service running and billing — and the operator's own `on_kill` doesn't cover that either, since a deferred task isn't on a worker.
The code itself is right and forward-compatible; it's the version gap that's worth surfacing. A sentence in the `deferrable` docstring or in `docs/operators/snowpark_containers.rst` saying that kill-time cleanup of a deferred job requires Airflow 3.3+ would save someone on 3.1 an unexplained orphaned service.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5248222242
@SameerMesiah97 I appreciate the review. I've addressed and replied to all of your comments. Let me know your thoughts. Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3754549622
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
Review Comment:
I agree the trigger isn't the most obvious place for these, but I'm not sure the hook is the best place either. The `SnowflakeHook` doesn't directly reference them at all. In dbt Cloud the hook actually references the statuses. Happy to move them if you feel strongly though.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -242,10 +247,35 @@ def execute(self, context: Context) -> str:
raise RuntimeError("Job name was not returned")
if not self.wait_for_completion:
return self.job_name
+ if self.deferrable:
+ self.defer(
+ trigger=SnowparkContainerJobTrigger(
+ job_name=self.job_name,
+ snowflake_conn_id=self.snowflake_conn_id,
+ poll_interval=self.poll_interval,
+ end_time=time.time() + self.timeout,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ warehouse=self.warehouse,
+ ),
+ # Pad past the trigger's end_time so its timeout event, which drops the service,
+ # fires before this hard backstop. A user-set execution_timeout takes precedence.
+ timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60),
Review Comment:
Apperciate the suggestion. I agree the execution deadline approach makes more sense here. I've updated the operator to pass an `execution_deadline` to the trigger and padded the `defer()` timeout so the trigger can emit its timeout event first. It now times out at whichever deadline (execution_deadline or end_time) is reached first. Added coverage for the new behavior as well.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3778937255
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
Review Comment:
I do believe the hook is at a lower-level of abstraction so any 'primitives' such as enums used my multiple modules naturally belong there. I would not say it is a serious issue but importing it like this from the trigger does look a bit strange.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3952909733
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -137,6 +117,8 @@ def __init__(
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,
Review Comment:
Good call. I've added a line to the change log.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890092549
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,12 +191,16 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ end_time = time.time() + self.timeout
while True:
+ if time.time() >= end_time:
+ self._drop_service()
Review Comment:
On timeout the service is dropped before any container logs are fetched, and a job that hung is exactly the case where those logs are wanted — dropping the service destroys them. Consider `self._log_container_output(status)` before the drop here and in the `"timeout"` branch of `execute_complete`. Separately: this path drops unconditionally, ignoring `drop_on_completion=False`. It's documented on the `timeout` param so it may well be deliberate, but it's worth a second look, since a *failed* job is left in place while a *timed-out* one is not.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
These Snowpark Container Services statuses now live in the SQL hook module, which is otherwise `SnowflakeHook` plus its OpenLineage surface. I take it the move is to break the operator↔trigger import cycle — `triggers/snowpark_containers.py` would do that equally well, since the operator already imports the trigger and the trigger imports nothing from the operator. Non-blocking, just a placement thought.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -242,10 +246,47 @@ def execute(self, context: Context) -> str:
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 = now + self.execution_timeout.total_seconds()
Review Comment:
`now` is `time.time()` at the moment of deferral, so everything before this point — connection setup, `_submit_job`, the `DESCRIBE` round trip — falls outside the window. `execution_timeout` is a task-level limit measured from task start, so a task with `execution_timeout=120` that spent 30s submitting gets 150s of wall clock. The comment two lines up says this "enforces the task-level limit", which is what makes it worth fixing rather than documenting. Deriving the deadline from the task's start instead would be exact:
```suggestion
execution_deadline = (
context["ti"].start_date.timestamp() + self.execution_timeout.total_seconds()
)
```
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -19,51 +19,25 @@
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.hooks.snowflake import (
+ CONTAINER_JOB_NON_TERMINAL_STATUSES,
Review Comment:
`TERMINAL_STATUSES` and `NON_TERMINAL_STATUSES` were public module-level names in this module and shipped in 6.15.0 (`Add SnowparkContainerJobOperator (#68259)`), still present in 6.16.1. Renaming them to `CONTAINER_JOB_*` and moving them to the hook turns `from airflow.providers.snowflake.operators.snowpark_containers import TERMINAL_STATUSES` into an `ImportError` on what is otherwise a feature release. `SnowparkContainerJobStatus` is unaffected — it's re-imported here — so it's just these two. Two aliases after the import block would keep the old spelling working.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5472622310
Thanks for the review. I've addressed all of your comments. Please see my response in regards to the placement of the statuses as I have made some changes there. I also re-tested this e2e against my Snowflake account, specifically the new changes which log the container outputs on timeout. It's working as expected.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5462891776
Would appreciate a review on this. Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890928483
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
While I was working on a follow up to implement retries (so a transient failure doesn't terminate the trigger/operator), I realized the SPCS trigger and operator will need some shared helpers beyond the statuses. So instead of putting them in the hook, I've moved them into a util new file in the existing `/utils` directory. This avoids the need to alias the status imports to maintain backwards compatibility and enables adding new helpers that are only used for SPCS.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on PR #70103:
URL: https://github.com/apache/airflow/pull/70103#issuecomment-5152892741
Would appreciate a review on this. Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3780157740
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
Review Comment:
Fair, I've moved them to the hook. I've prefixed the status sets with CONTAINER_JOB_* since bare TERMINAL_STATUSES was now ambiguous about what it belonged to.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -248,20 +247,29 @@ def execute(self, context: Context) -> str:
if not self.wait_for_completion:
return self.job_name
if self.deferrable:
Review Comment:
Added a comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3754550597
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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:
+ 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
+
+ if time.time() > self.end_time:
Review Comment:
I moved timeout checks above the job status evaluation so that the hard deadline is respected. I also updated the timeout check in the operator to align the two.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3754550597
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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:
+ 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
+
+ if time.time() > self.end_time:
Review Comment:
I moved timeout checks above the job status evaluation so that the hard deadline is respected. Note - I also updated the timeout check in the operator to align the two.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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:
+ try:
+ status = await self._describe_status(hook=hook)
+ except Exception as e:
Review Comment:
I actually originally implemented a max consecutive failures approach here but I removed it from this PR as I felt that change was better suited for a separate PR. I think it should apply to the operator's polling too and not just the trigger. Planning on doing this in a follow up.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,179 @@
+# 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 enum import Enum
+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.triggers.base import BaseTrigger, TriggerEvent
+
+
+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,
+ }
+)
+
+
+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 database: (Optional) name of database.
+ :param schema: (Optional) name of schema.
+ :param role: (Optional) name of role.
+ :param warehouse: (Optional) name of warehouse.
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ 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.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,
+ "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
Review Comment:
Based on their documentation, yes.
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowpark_containers.py:
##########
@@ -280,3 +312,58 @@ def test_execute_skips_drop_when_disabled(self, mock_hook_cls, mock_submit, mock
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"
Review Comment:
Yep, I added coverage for the new deadline handling.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to SnowparkContainerJobOperator [airflow]
Posted by "justinpakzad (via GitHub)" <gi...@apache.org>.
justinpakzad commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3952909174
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,7 +190,14 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ status = None
+ end_time = time.time() + self.timeout
Review Comment:
Good call, I've updated `_poll_for_status` to use `time.monotonic()`. Also updated one of the tests since it was using time machine which doesn't seem to mock `time.monotonic()`
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -137,6 +117,8 @@ def __init__(
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,
Review Comment:
Good call. I've added a line to the change log. I also squashed the commits (I know maintainers do it anyways but the commit message weren't the best) and added a more detailed message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org