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