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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
