Source code for

# 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
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# 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 warnings
from functools import cached_property
from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence

from import Operation

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
from import (
from airflow.triggers.base import BaseTrigger, TriggerEvent

    from datetime import datetime

    from kubernetes_asyncio.client import V1Job

[docs]class GKEStartPodTrigger(KubernetesPodTrigger): """ Trigger for checking pod status until it finishes its job. :param pod_name: The name of the pod. :param pod_namespace: The namespace of the pod. :param cluster_url: The URL pointed to the cluster. :param ssl_ca_cert: SSL certificate that is used for authentication to the pod. :param cluster_context: Context that points to kubernetes cluster. :param poll_interval: Polling period in seconds to check for the status. :param trigger_start_time: time in Datetime format when the trigger was started :param in_cluster: run kubernetes client with in_cluster configuration. :param get_logs: get the stdout of the container as logs of the tasks. :param startup_timeout: timeout in seconds to start up the pod. :param base_container_name: The name of the base container in the pod. This container's logs will appear as part of this task's logs if get_logs is True. Defaults to None. If None, will consult the class variable BASE_CONTAINER_NAME (which defaults to "base") for the base container name to use. :param on_finish_action: What to do when the pod reaches its final state, or the execution is interrupted. If "delete_pod", the pod will be deleted regardless its state; if "delete_succeeded_pod", only succeeded pod will be deleted. You can set to "keep_pod" to keep the pod. :param should_delete_pod: What to do when the pod reaches its final state, or the execution is interrupted. If True (default), delete the pod; if False, leave the pod. Deprecated - use `on_finish_action` instead. """ def __init__( self, pod_name: str, pod_namespace: str, cluster_url: str, ssl_ca_cert: str, base_container_name: str, trigger_start_time: datetime, cluster_context: str | None = None, poll_interval: float = 2, in_cluster: bool | None = None, get_logs: bool = True, startup_timeout: int = 120, on_finish_action: str = "delete_pod", should_delete_pod: bool | None = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, *args, **kwargs, ): super().__init__( pod_name, pod_namespace, trigger_start_time, base_container_name, *args, **kwargs, ) self.pod_name = pod_name self.pod_namespace = pod_namespace self.trigger_start_time = trigger_start_time self.base_container_name = base_container_name self.poll_interval = poll_interval self.cluster_context = cluster_context self.in_cluster = in_cluster self.get_logs = get_logs self.startup_timeout = startup_timeout self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain if should_delete_pod is not None: warnings.warn( "`should_delete_pod` parameter is deprecated, please use `on_finish_action`", AirflowProviderDeprecationWarning, stacklevel=2, ) self.on_finish_action = ( OnFinishAction.DELETE_POD if should_delete_pod else OnFinishAction.KEEP_POD ) self.should_delete_pod = should_delete_pod else: self.on_finish_action = OnFinishAction(on_finish_action) self.should_delete_pod = self.on_finish_action == OnFinishAction.DELETE_POD self._cluster_url = cluster_url self._ssl_ca_cert = ssl_ca_cert
[docs] def serialize(self) -> tuple[str, dict[str, Any]]: return ( "", { "pod_name": self.pod_name, "pod_namespace": self.pod_namespace, "cluster_url": self._cluster_url, "ssl_ca_cert": self._ssl_ca_cert, "poll_interval": self.poll_interval, "cluster_context": self.cluster_context, "in_cluster": self.in_cluster, "get_logs": self.get_logs, "startup_timeout": self.startup_timeout, "trigger_start_time": self.trigger_start_time, "base_container_name": self.base_container_name, "should_delete_pod": self.should_delete_pod, "on_finish_action": self.on_finish_action.value, "gcp_conn_id": self.gcp_conn_id, "impersonation_chain": self.impersonation_chain, }, )
[docs] def hook(self) -> GKEPodAsyncHook: # type: ignore[override] return GKEPodAsyncHook( cluster_url=self._cluster_url, ssl_ca_cert=self._ssl_ca_cert, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, enable_tcp_keepalive=True, )
[docs]class GKEOperationTrigger(BaseTrigger): """Trigger which checks status of the operation.""" def __init__( self, operation_name: str, project_id: str | None, location: str, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, poll_interval: int = 10, ): super().__init__() self.operation_name = operation_name self.project_id = project_id self.location = location self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain self.poll_interval = poll_interval self._hook: GKEAsyncHook | None = None
[docs] def serialize(self) -> tuple[str, dict[str, Any]]: """Serialize GKEOperationTrigger arguments and classpath.""" return ( "", { "operation_name": self.operation_name, "project_id": self.project_id, "location": self.location, "gcp_conn_id": self.gcp_conn_id, "impersonation_chain": self.impersonation_chain, "poll_interval": self.poll_interval, }, )
[docs] async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override] """Get operation status and yields corresponding event.""" hook = self._get_hook() try: while True: operation = await hook.get_operation( operation_name=self.operation_name, project_id=self.project_id, ) status = operation.status if status == Operation.Status.DONE: yield TriggerEvent( { "status": "success", "message": "Operation is successfully ended.", "operation_name":, } ) return elif status in (Operation.Status.RUNNING, Operation.Status.PENDING):"Operation is still running.")"Sleeping for %ss...", self.poll_interval) await asyncio.sleep(self.poll_interval) else: yield TriggerEvent( { "status": "failed", "message": f"Operation has failed with status: {operation.status}", } ) return except Exception as e: self.log.exception("Exception occurred while checking operation status") yield TriggerEvent( { "status": "error", "message": str(e), } )
def _get_hook(self) -> GKEAsyncHook: if self._hook is None: self._hook = GKEAsyncHook( gcp_conn_id=self.gcp_conn_id, location=self.location, impersonation_chain=self.impersonation_chain, ) return self._hook
[docs]class GKEJobTrigger(BaseTrigger): """GKEJobTrigger run on the trigger worker to check the state of Job.""" def __init__( self, cluster_url: str, ssl_ca_cert: str, job_name: str, job_namespace: str, gcp_conn_id: str = "google_cloud_default", poll_interval: float = 2, impersonation_chain: str | Sequence[str] | None = None, ) -> None: super().__init__() self.cluster_url = cluster_url self.ssl_ca_cert = ssl_ca_cert self.job_name = job_name self.job_namespace = job_namespace self.gcp_conn_id = gcp_conn_id self.poll_interval = poll_interval self.impersonation_chain = impersonation_chain
[docs] def serialize(self) -> tuple[str, dict[str, Any]]: """Serialize KubernetesCreateJobTrigger arguments and classpath.""" return ( "", { "cluster_url": self.cluster_url, "ssl_ca_cert": self.ssl_ca_cert, "job_name": self.job_name, "job_namespace": self.job_namespace, "gcp_conn_id": self.gcp_conn_id, "poll_interval": self.poll_interval, "impersonation_chain": self.impersonation_chain, }, )
[docs] async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override] """Get current job status and yield a TriggerEvent.""" job: V1Job = await self.hook.wait_until_job_complete(name=self.job_name, namespace=self.job_namespace) job_dict = job.to_dict() error_message = self.hook.is_job_failed(job=job) status = "error" if error_message else "success" message = f"Job failed with error: {error_message}" if error_message else "Job completed successfully" yield TriggerEvent( { "name":, "namespace": job.metadata.namespace, "status": status, "message": message, "job": job_dict, } )
[docs] def hook(self) -> GKEKubernetesAsyncHook: return GKEKubernetesAsyncHook( cluster_url=self.cluster_url, ssl_ca_cert=self.ssl_ca_cert, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, )

Was this entry helpful?