# 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 warnings
from functools import cached_property
from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence
from google.cloud.container_v1.types 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 airflow.providers.google.cloud.hooks.kubernetes_engine import (
GKEAsyncHook,
GKEKubernetesAsyncHook,
)
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
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 (
"airflow.providers.google.cloud.triggers.kubernetes_engine.GKEStartPodTrigger",
{
"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,
"logging_interval": self.logging_interval,
"last_log_time": self.last_log_time,
},
)
@cached_property
[docs] def hook(self) -> GKEKubernetesAsyncHook: # type: ignore[override]
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,
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 (
"airflow.providers.google.cloud.triggers.kubernetes_engine.GKEOperationTrigger",
{
"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": operation.name,
}
)
return
elif status in (Operation.Status.RUNNING, Operation.Status.PENDING):
self.log.info("Operation is still running.")
self.log.info("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 (
"airflow.providers.google.cloud.triggers.kubernetes_engine.GKEJobTrigger",
{
"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": job.metadata.name,
"namespace": job.metadata.namespace,
"status": status,
"message": message,
"job": job_dict,
}
)
@cached_property
[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,
)