From cfcd75213e5c4fd1b983da37fb5a6c2e320edc1f Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Tue, 31 Jan 2023 14:55:57 +0100 Subject: [PATCH 1/5] Add deferrable mode to GKEStartPodOperator --- .../google/cloud/hooks/kubernetes_engine.py | 209 +++++++++++++- .../cloud/operators/kubernetes_engine.py | 191 ++++++------- .../cloud/triggers/kubernetes_engine.py | 97 ++++++- .../operators/cloud/kubernetes_engine.rst | 8 + .../cloud/hooks/test_kubernetes_engine.py | 73 ++++- .../cloud/operators/test_kubernetes_engine.py | 251 +++++------------ .../cloud/triggers/test_kubernetes_engine.py | 256 +++++++++++++++++- .../example_kubernetes_engine_async.py | 12 +- 8 files changed, 805 insertions(+), 292 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 838f547f76e26..c5645623ef497 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -25,22 +25,33 @@ """ from __future__ import annotations +import contextlib import json import time import warnings from typing import Sequence -from google.api_core.exceptions import NotFound +import google.auth.credentials +from gcloud.aio.auth import Token +from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry +from google.auth.transport import requests as google_requests # not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable from google.cloud import container_v1, exceptions # type: ignore[attr-defined] from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient from google.cloud.container_v1.types import Cluster, Operation +from kubernetes import client +from kubernetes_asyncio import client as async_client +from kubernetes_asyncio.client.models import V1Pod +from kubernetes_asyncio.config.kube_config import FileOrData +from urllib3.exceptions import HTTPError from airflow import version +from airflow.compat.functools import cached_property from airflow.exceptions import AirflowException +from airflow.kubernetes.pod_generator_deprecated import PodDefaults from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import ( PROVIDE_PROJECT_ID, @@ -66,10 +77,6 @@ def __init__( location: str | None = None, impersonation_chain: str | Sequence[str] | None = None, ) -> None: - if delegate_to: - warnings.warn( - "'delegate_to' parameter is deprecated, please use 'impersonation_chain'", DeprecationWarning - ) super().__init__( gcp_conn_id=gcp_conn_id, delegate_to=delegate_to, @@ -336,3 +343,195 @@ async def get_operation( return await client.get_operation( name=operation_path, ) + + +class GKEPodHook(GoogleBaseHook): + """Hook for managing pods.""" + + def __init__( + self, + cluster_url: str, + ssl_ca_cert: str, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self._cluster_url = cluster_url + self._ssl_ca_cert = ssl_ca_cert + + @cached_property + def api_client(self) -> client.ApiClient: + return self.get_conn() + + @cached_property + def core_v1_client(self) -> client.CoreV1Api: + return client.CoreV1Api(self.api_client) + + @property + def is_in_cluster(self) -> bool: + return False + + @staticmethod + def get_xcom_sidecar_container_image(): + """Returns the xcom sidecar image that defined in the connection""" + return PodDefaults.SIDECAR_CONTAINER.image + + def get_conn(self) -> client.ApiClient: + configuration = self._get_config() + return client.ApiClient(configuration) + + def _get_config(self) -> client.configuration.Configuration: + configuration = client.Configuration( + host=self._cluster_url, + api_key_prefix={"authorization": "Bearer"}, + api_key={"authorization": self._get_token(self.get_credentials())}, + ) + configuration.ssl_ca_cert = FileOrData( + { + "certificate-authority-data": self._ssl_ca_cert, + }, + file_key_name="certificate-authority", + ).as_file() + return configuration + + @staticmethod + def _get_token(creds: google.auth.credentials.Credentials) -> str: + if creds.token is None or creds.expired: + auth_req = google_requests.Request() + creds.refresh(auth_req) + return creds.token + + def get_pod(self, name: str, namespace: str) -> V1Pod: + """ + Gets pod's object. + + :param name: Name of the pod. + :param namespace: Name of the pod's namespace. + """ + return self.core_v1_client.read_namespaced_pod( + name=name, + namespace=namespace, + ) + + +class AsyncGKEPodHook(GoogleBaseAsyncHook): + """ + Hook for managing GKE pods in asynchronous way. + + :param cluster_url: The URL pointed to the cluster. + :param ssl_ca_cert: SSL certificate that is used for authentication to the pod. + """ + + sync_hook_class = GKEPodHook + scopes = ["https://www.googleapis.com/auth/cloud-platform"] + + def __init__( + self, + cluster_url: str, + ssl_ca_cert: str, + **kwargs, + ): + + self._cluster_url = cluster_url + self._ssl_ca_cert = ssl_ca_cert + + kwargs.update( + cluster_url=cluster_url, + ssl_ca_cert=ssl_ca_cert, + ) + super().__init__(**kwargs) + + @contextlib.asynccontextmanager + async def get_conn(self, token: Token) -> async_client.ApiClient: # type: ignore[override] + kube_client = None + try: + kube_client = await self._load_config(token) + yield kube_client + finally: + if kube_client is not None: + await kube_client.close() + + async def _load_config(self, token: Token) -> async_client.ApiClient: + configuration = self._get_config() + access_token = await token.get() + return async_client.ApiClient( + configuration, + header_name="Authorization", + header_value=f"Bearer {access_token}", + ) + + def _get_config(self) -> async_client.configuration.Configuration: + configuration = async_client.Configuration( + host=self._cluster_url, + ssl_ca_cert=FileOrData( + { + "certificate-authority-data": self._ssl_ca_cert, + }, + file_key_name="certificate-authority", + ).as_file(), + ) + return configuration + + async def get_pod(self, name: str, namespace: str) -> V1Pod: + """ + Gets pod's object. + + :param name: Name of the pod. + :param namespace: Name of the pod's namespace. + """ + async with Token(scopes=self.scopes) as token: + async with self.get_conn(token) as connection: + v1_api = async_client.CoreV1Api(connection) + pod: V1Pod = await v1_api.read_namespaced_pod( + name=name, + namespace=namespace, + ) + return pod + + async def delete_pod(self, name: str, namespace: str): + """ + Deletes pod's object. + + :param name: Name of the pod. + :param namespace: Name of the pod's namespace. + """ + async with Token(scopes=self.scopes) as token: + async with self.get_conn(token) as connection: + try: + v1_api = async_client.CoreV1Api(connection) + await v1_api.delete_namespaced_pod( + name=name, + namespace=namespace, + body=client.V1DeleteOptions(), + ) + except async_client.ApiException as e: + # If the pod is already deleted + if e.status != 404: + raise + + async def read_logs(self, name: str, namespace: str): + """ + Reads logs inside the pod while starting containers inside. All the logs will be outputted with its + timestamp to track the logs after the execution of the pod is completed. The method is used for async + output of the logs only in the pod failed it execution or the task was cancelled by the user. + + :param name: Name of the pod. + :param namespace: Name of the pod's namespace. + """ + async with Token(scopes=self.scopes) as token: + async with self.get_conn(token) as connection: + try: + v1_api = async_client.CoreV1Api(connection) + logs = await v1_api.read_namespaced_pod_log( + name=name, + namespace=namespace, + follow=False, + timestamps=True, + ) + logs = logs.splitlines() + for line in logs: + self.log.info("Container logs from %s", line) + return logs + except HTTPError: + self.log.exception("There was an error reading the kubernetes API.") + raise diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index 9482f95dd443c..84466f27adbe8 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -18,15 +18,14 @@ """This module contains Google Kubernetes Engine operators.""" from __future__ import annotations -import os -import tempfile import warnings -from contextlib import contextmanager -from typing import TYPE_CHECKING, Generator, Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.exceptions import AlreadyExists from google.cloud.container_v1.types import Cluster +from kubernetes.client.models import V1Pod +from airflow.compat.functools import cached_property from airflow.exceptions import AirflowException try: @@ -34,15 +33,14 @@ except ImportError: # preserve backward compatibility for older versions of cncf.kubernetes provider from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator -from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEHook +from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEHook, GKEPodHook from airflow.providers.google.cloud.links.kubernetes_engine import ( KubernetesEngineClusterLink, KubernetesEnginePodLink, ) from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger -from airflow.providers.google.common.hooks.base_google import GoogleBaseHook -from airflow.utils.process_utils import execute_in_subprocess, patch_environ +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEPodTrigger +from airflow.utils.timezone import utcnow if TYPE_CHECKING: from airflow.utils.context import Context @@ -233,6 +231,8 @@ class GKECreateClusterOperator(GoogleCloudBaseOperator): If set as a sequence, the identities from the list must grant Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated). + :param deferrable: Run operator in the deferrable mode. + :param poll_interval: Interval size which defines how often operation status is checked. """ template_fields: Sequence[str] = ( @@ -363,7 +363,7 @@ def _get_hook(self) -> GKEHook: class GKEStartPodOperator(KubernetesPodOperator): """ Executes a task in a Kubernetes pod in the specified Google Kubernetes - Engine cluster + Engine cluster. This Operator assumes that the system has gcloud installed and has configured a connection id with a service account. @@ -386,7 +386,7 @@ class GKEStartPodOperator(KubernetesPodOperator): should be spawned in :param use_internal_ip: Use the internal IP address as the endpoint. :param project_id: The Google Developers Console project id - :param gcp_conn_id: The google cloud connection id to use. This allows for + :param gcp_conn_id: The Google cloud connection id to use. This allows for users to specify a service account. :param impersonation_chain: Optional service account to impersonate using short-term credentials, or list of accounts required to get the access_token @@ -399,8 +399,9 @@ class GKEStartPodOperator(KubernetesPodOperator): :param regional: The location param is region name. :param is_delete_operator_pod: What to do when the pod reaches its final state, or the execution is interrupted. If True, delete the - pod; if False, leave the pod. Current default is False, but this will be + pod; if False, leave the pod. Current default is False, but this will be changed in the next major release of this provider. + :param deferrable: Run operator in the deferrable mode. """ template_fields: Sequence[str] = tuple( @@ -413,11 +414,11 @@ def __init__( *, location: str, cluster_name: str, - use_internal_ip: bool = False, + use_internal_ip: bool | None = None, project_id: str | None = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, - regional: bool = False, + regional: bool | None = None, is_delete_operator_pod: bool | None = None, **kwargs, ) -> None: @@ -432,14 +433,34 @@ def __init__( ) is_delete_operator_pod = False + if use_internal_ip is not None: + warnings.warn( + f"You have set parameter use_internal_ip in class {self.__class__.__name__}. " + "In current implementation of the operator the parameter is not used and will " + "be deleted in future.", + DeprecationWarning, + stacklevel=2, + ) + + if regional is not None: + warnings.warn( + f"You have set parameter regional in class {self.__class__.__name__}. " + "In current implementation of the operator the parameter is not used and will " + "be deleted in future.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(is_delete_operator_pod=is_delete_operator_pod, **kwargs) self.project_id = project_id self.location = location self.cluster_name = cluster_name self.gcp_conn_id = gcp_conn_id - self.use_internal_ip = use_internal_ip self.impersonation_chain = impersonation_chain - self.regional = regional + + self.pod: V1Pod | None = None + self._ssl_ca_cert: str | None = None + self._cluster_url: str | None = None if self.gcp_conn_id is None: raise AirflowException( @@ -452,89 +473,69 @@ def __init__( if self.config_file: raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.") - @staticmethod - @contextmanager - def get_gke_config_file( - gcp_conn_id, - project_id: str | None, - cluster_name: str, - impersonation_chain: str | Sequence[str] | None, - regional: bool, - location: str, - use_internal_ip: bool, - ) -> Generator[str, None, None]: - - hook = GoogleBaseHook(gcp_conn_id=gcp_conn_id) - project_id = project_id or hook.project_id + @cached_property + def cluster_hook(self) -> GKEHook: + return GKEHook( + gcp_conn_id=self.gcp_conn_id, + location=self.location, + impersonation_chain=self.impersonation_chain, + ) - if not project_id: - raise AirflowException( - "The project id must be passed either as " - "keyword project_id parameter or as project_id extra " - "in Google Cloud connection definition. Both are not set!" + @cached_property + def hook(self) -> GKEPodHook: # type: ignore[override] + if self._cluster_url is None or self._ssl_ca_cert is None: + raise AttributeError( + "Cluster url and ssl_ca_cert should be defined before using self.hook method. " + "Try to use self.get_kube_creds method", ) - # Write config to a temp file and set the environment variable to point to it. - # This is to avoid race conditions of reading/writing a single file - with tempfile.NamedTemporaryFile() as conf_file, patch_environ( - {KUBE_CONFIG_ENV_VAR: conf_file.name} - ), hook.provide_authorized_gcloud(): - # Attempt to get/update credentials - # We call gcloud directly instead of using google-cloud-python api - # because there is no way to write kubernetes config to a file, which is - # required by KubernetesPodOperator. - # The gcloud command looks at the env variable `KUBECONFIG` for where to save - # the kubernetes config file. - cmd = [ - "gcloud", - "container", - "clusters", - "get-credentials", - cluster_name, - "--project", - project_id, - ] - if impersonation_chain: - if isinstance(impersonation_chain, str): - impersonation_account = impersonation_chain - elif len(impersonation_chain) == 1: - impersonation_account = impersonation_chain[0] - else: - raise AirflowException( - "Chained list of accounts is not supported, please specify only one service account" - ) - - cmd.extend( - [ - "--impersonate-service-account", - impersonation_account, - ] - ) - if regional: - cmd.append("--region") - else: - cmd.append("--zone") - cmd.append(location) - if use_internal_ip: - cmd.append("--internal-ip") - execute_in_subprocess(cmd) - - # Tell `KubernetesPodOperator` where the config file is located - yield os.environ[KUBE_CONFIG_ENV_VAR] + hook = GKEPodHook( + cluster_url=self._cluster_url, + ssl_ca_cert=self._ssl_ca_cert, + ) + return hook - def execute(self, context: Context) -> str | None: + def execute(self, context: Context): + """Executes process of creating pod and executing provided command inside it.""" + self.fetch_cluster_info() + return super().execute(context) - with GKEStartPodOperator.get_gke_config_file( - gcp_conn_id=self.gcp_conn_id, + def fetch_cluster_info(self) -> tuple[str, str | None]: + """Fetches cluster info for connecting to it.""" + cluster = self.cluster_hook.get_cluster( + name=self.cluster_name, project_id=self.project_id, - cluster_name=self.cluster_name, - impersonation_chain=self.impersonation_chain, - regional=self.regional, - location=self.location, - use_internal_ip=self.use_internal_ip, - ) as config_file: - self.config_file = config_file - result = super().execute(context) - if not self.is_delete_operator_pod: - KubernetesEnginePodLink.persist(context=context, task_instance=self) - return result + ) + + self._cluster_url = f"https://{cluster.endpoint}" + self._ssl_ca_cert = cluster.master_auth.cluster_ca_certificate + return self._cluster_url, self._ssl_ca_cert + + def invoke_defer_method(self): + """Method to easily redefine triggers which are being used in child classes.""" + trigger_start_time = utcnow() + self.defer( + trigger=GKEPodTrigger( + pod_name=self.pod.metadata.name, + pod_namespace=self.pod.metadata.namespace, + trigger_start_time=trigger_start_time, + cluster_url=self._cluster_url, + ssl_ca_cert=self._ssl_ca_cert, + get_logs=self.get_logs, + startup_timeout=self.startup_timeout_seconds, + cluster_context=self.cluster_context, + poll_interval=self.poll_interval, + in_cluster=self.in_cluster, + should_delete_pod=self.is_delete_operator_pod, + base_container_name=self.base_container_name, + ), + method_name="execute_complete", + kwargs={"cluster_url": self._cluster_url, "ssl_ca_cert": self._ssl_ca_cert}, + ) + + def execute_complete(self, context: Context, event: dict, **kwargs): + # It is required for hook to be initialized + self._cluster_url = kwargs["cluster_url"] + self._ssl_ca_cert = kwargs["ssl_ca_cert"] + + return super().execute_complete(context, event, **kwargs) diff --git a/airflow/providers/google/cloud/triggers/kubernetes_engine.py b/airflow/providers/google/cloud/triggers/kubernetes_engine.py index 9d446bbc41c32..5508abc9975fc 100644 --- a/airflow/providers/google/cloud/triggers/kubernetes_engine.py +++ b/airflow/providers/google/cloud/triggers/kubernetes_engine.py @@ -18,12 +18,106 @@ from __future__ import annotations import asyncio +from datetime import datetime from typing import Any, AsyncIterator, Sequence from google.cloud.container_v1.types import Operation -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook +try: + from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger +except ImportError: + # preserve backward compatibility for older versions of cncf.kubernetes provider + from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import KubernetesPodTrigger from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEPodHook, AsyncGKEHook + + +class GKEPodTrigger(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 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. + :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. + """ + + 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, + should_delete_pod: bool = True, + get_logs: bool = True, + startup_timeout: int = 120, + *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.should_delete_pod = should_delete_pod + self.get_logs = get_logs + self.startup_timeout = startup_timeout + + self._cluster_url = cluster_url + self._ssl_ca_cert = ssl_ca_cert + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEPodTrigger", + { + "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, + "should_delete_pod": self.should_delete_pod, + "get_logs": self.get_logs, + "startup_timeout": self.startup_timeout, + "trigger_start_time": self.trigger_start_time, + "base_container_name": self.base_container_name, + }, + ) + + def _get_async_hook(self) -> AsyncGKEPodHook: # type: ignore[override] + return AsyncGKEPodHook( + cluster_url=self._cluster_url, + ssl_ca_cert=self._ssl_ca_cert, + ) class GKEOperationTrigger(BaseTrigger): @@ -119,3 +213,4 @@ def _get_hook(self) -> AsyncGKEHook: impersonation_chain=self.impersonation_chain, ) return self._hook + diff --git a/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst b/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst index 8bcadef6bdf31..b0545a5b3d7d1 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst @@ -148,6 +148,14 @@ And then use it in other operators: :start-after: [START howto_operator_gke_xcom_result] :end-before: [END howto_operator_gke_xcom_result] +You can use deferrable mode for this action in order to run the operator asynchronously: + +.. exampleinclude:: /../../tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_gke_start_pod_xcom_async] + :end-before: [END howto_operator_gke_start_pod_xcom_async] + Reference ^^^^^^^^^ diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index 806814996c6ee..da1a31aaea9ad 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -20,12 +20,13 @@ import sys from asyncio import Future +import kubernetes.client import pytest from google.cloud.container_v1 import ClusterManagerAsyncClient from google.cloud.container_v1.types import Cluster from airflow.exceptions import AirflowException -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, GKEHook +from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, GKEHook, AsyncGKEPodHook from airflow.providers.google.common.consts import CLIENT_INFO from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id @@ -40,6 +41,10 @@ GKE_ZONE = "test-zone" BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" GKE_STRING = "airflow.providers.google.cloud.hooks.kubernetes_engine.{}" +CLUSTER_URL = "https://path.to.cluster" +SSL_CA_CERT = "test-ssl-ca-cert" +POD_NAME = "test-pod-name" +POD_NAMESPACE = "test" ASYNC_HOOK_STRING = GKE_STRING.format("AsyncGKEHook") GCP_CONN_ID = "test-gcp-conn-id" DELEGATE_TO = "test-delegate-to" @@ -291,6 +296,72 @@ def test_wait_for_response_running(self, time_mock, operation_mock): assert operation_mock.call_count == 2 +class TestAsyncGKEPodHook: + @staticmethod + def make_mock_awaitable(mock_obj, result=None): + f = Future() + f.set_result(result) + + mock_obj.return_value = f + + @pytest.fixture() + def async_hook(self): + return AsyncGKEPodHook( + cluster_url=CLUSTER_URL, + ssl_ca_cert=SSL_CA_CERT, + ) + + @pytest.mark.asyncio + @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod")) + async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook): + self.make_mock_awaitable(read_namespace_pod_mock) + + await async_hook.get_pod(name=POD_NAME, namespace=POD_NAMESPACE) + + get_conn_mock.assert_called_once() + read_namespace_pod_mock.assert_called_with( + name=POD_NAME, + namespace=POD_NAMESPACE, + ) + + @pytest.mark.asyncio + @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("async_client.CoreV1Api.delete_namespaced_pod")) + async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook): + self.make_mock_awaitable(delete_namespaced_pod) + + await async_hook.delete_pod(name=POD_NAME, namespace=POD_NAMESPACE) + + get_conn_mock.assert_called_once() + delete_namespaced_pod.assert_called_with( + name=POD_NAME, + namespace=POD_NAMESPACE, + body=kubernetes.client.V1DeleteOptions(), + ) + + @pytest.mark.asyncio + @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) + @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod_log")) + async def test_read_logs(self, read_namespaced_pod_log, get_conn_mock, async_hook, caplog): + self.make_mock_awaitable(read_namespaced_pod_log, result="Test string #1\nTest string #2\n") + + await async_hook.read_logs(name=POD_NAME, namespace=POD_NAMESPACE) + + get_conn_mock.assert_called_once() + read_namespaced_pod_log.assert_called_with( + name=POD_NAME, + namespace=POD_NAMESPACE, + follow=False, + timestamps=True, + ) + assert "Test string #1" in caplog.text + assert "Test string #2" in caplog.text + + @pytest.fixture() def async_gke_hook(): return AsyncGKEHook( diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index 5892e01a5a5b5..abc006b6d9d86 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -20,11 +20,10 @@ import json import os from unittest import mock -from unittest.mock import PropertyMock import pytest -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import Connection from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator from airflow.providers.google.cloud.operators.kubernetes_engine import ( @@ -32,6 +31,7 @@ GKEDeleteClusterOperator, GKEStartPodOperator, ) +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEPodTrigger TEST_GCP_PROJECT_ID = "test-id" PROJECT_LOCATION = "test-location" @@ -59,6 +59,14 @@ GCLOUD_COMMAND = "gcloud container clusters get-credentials {} --zone {} --project {}" KUBE_ENV_VAR = "KUBECONFIG" FILE_NAME = "/tmp/mock_name" +KUB_OP_PATH = "airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.{}" +GKE_HOOK_MODULE_PATH = "airflow.providers.google.cloud.operators.kubernetes_engine" +GKE_HOOK_PATH = f"{GKE_HOOK_MODULE_PATH}.GKEClusterHook" +KUB_OPERATOR_EXEC = "airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute" +TEMP_FILE = "tempfile.NamedTemporaryFile" +GKE_OP_PATH = "airflow.providers.google.cloud.operators.kubernetes_engine.GKEStartPodOperator" +CLUSTER_URL = "https://test-host" +SSL_CA_CERT = "TEST_SSL_CA_CERT_CONTENT" class TestGoogleCloudPlatformContainerOperator: @@ -71,7 +79,7 @@ class TestGoogleCloudPlatformContainerOperator: PROJECT_BODY_CREATE_CLUSTER_NODE_POOLS, ], ) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_create_execute(self, mock_hook, body): operator = GKECreateClusterOperator( project_id=TEST_GCP_PROJECT_ID, location=PROJECT_LOCATION, body=body, task_id=PROJECT_TASK_ID @@ -129,19 +137,19 @@ def test_create_execute(self, mock_hook, body): )(), ], ) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_create_execute_error_body(self, mock_hook, body): with pytest.raises(AirflowException): GKECreateClusterOperator( project_id=TEST_GCP_PROJECT_ID, location=PROJECT_LOCATION, body=body, task_id=PROJECT_TASK_ID ) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_create_execute_error_project_id(self, mock_hook): with pytest.raises(AirflowException): GKECreateClusterOperator(location=PROJECT_LOCATION, body=PROJECT_BODY, task_id=PROJECT_TASK_ID) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_create_execute_error_location(self, mock_hook): with pytest.raises(AirflowException): GKECreateClusterOperator( @@ -179,19 +187,19 @@ def test_delete_execute(self, mock_hook): wait_to_complete=True, ) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_delete_execute_error_project_id(self, mock_hook): with pytest.raises(AirflowException): GKEDeleteClusterOperator(location=PROJECT_LOCATION, name=CLUSTER_NAME, task_id=PROJECT_TASK_ID) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_delete_execute_error_cluster_name(self, mock_hook): with pytest.raises(AirflowException): GKEDeleteClusterOperator( project_id=TEST_GCP_PROJECT_ID, location=PROJECT_LOCATION, task_id=PROJECT_TASK_ID ) - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GKEHook") + @mock.patch(GKE_HOOK_PATH) def test_delete_execute_error_location(self, mock_hook): with pytest.raises(AirflowException): GKEDeleteClusterOperator( @@ -234,75 +242,12 @@ def test_template_fields(self): assert set(KubernetesPodOperator.template_fields).issubset(GKEStartPodOperator.template_fields) @mock.patch.dict(os.environ, {}) - @mock.patch( - "airflow.hooks.base.BaseHook.get_connections", - return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], - ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") - def test_execute(self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock): - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) - + @mock.patch(KUB_OPERATOR_EXEC) + @mock.patch(TEMP_FILE) + @mock.patch(f"{GKE_OP_PATH}.fetch_cluster_info") + def test_execute(self, fetch_cluster_info_mock, file_mock, exec_mock): self.gke_op.execute(context=mock.MagicMock()) - - mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() - - mock_execute_in_subprocess.assert_called_once_with( - [ - "gcloud", - "container", - "clusters", - "get-credentials", - CLUSTER_NAME, - "--project", - TEST_GCP_PROJECT_ID, - "--zone", - PROJECT_LOCATION, - ] - ) - - assert self.gke_op.config_file == FILE_NAME - - @mock.patch.dict(os.environ, {}) - @mock.patch( - "airflow.hooks.base.BaseHook.get_connections", - return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], - ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") - def test_execute_regional( - self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock - ): - self.gke_op.regional = True - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) - - self.gke_op.execute(context=mock.MagicMock()) - - mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() - - mock_execute_in_subprocess.assert_called_once_with( - [ - "gcloud", - "container", - "clusters", - "get-credentials", - CLUSTER_NAME, - "--project", - TEST_GCP_PROJECT_ID, - "--region", - PROJECT_LOCATION, - ] - ) - - assert self.gke_op.config_file == FILE_NAME + fetch_cluster_info_mock.assert_called_once() def test_config_file_throws_error(self): with pytest.raises(AirflowException): @@ -322,133 +267,71 @@ def test_config_file_throws_error(self): "airflow.hooks.base.BaseHook.get_connections", return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") - def test_execute_with_internal_ip( - self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock - ): - self.gke_op.use_internal_ip = True - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) - - self.gke_op.execute(context=mock.MagicMock()) - - mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() - - mock_execute_in_subprocess.assert_called_once_with( - [ - "gcloud", - "container", - "clusters", - "get-credentials", - CLUSTER_NAME, - "--project", - TEST_GCP_PROJECT_ID, - "--zone", - PROJECT_LOCATION, - "--internal-ip", - ] - ) - - assert self.gke_op.config_file == FILE_NAME - - @mock.patch.dict(os.environ, {}) - @mock.patch( - "airflow.hooks.base.BaseHook.get_connections", - return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], - ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") + @mock.patch(KUB_OPERATOR_EXEC) + @mock.patch(TEMP_FILE) + @mock.patch(f"{GKE_OP_PATH}.fetch_cluster_info") def test_execute_with_impersonation_service_account( - self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock + self, fetch_cluster_info_mock, file_mock, exec_mock, get_con_mock ): - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) self.gke_op.impersonation_chain = "test_account@example.com" self.gke_op.execute(context=mock.MagicMock()) - - mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() - - mock_execute_in_subprocess.assert_called_once_with( - [ - "gcloud", - "container", - "clusters", - "get-credentials", - CLUSTER_NAME, - "--project", - TEST_GCP_PROJECT_ID, - "--impersonate-service-account", - "test_account@example.com", - "--zone", - PROJECT_LOCATION, - ] - ) - - assert self.gke_op.config_file == FILE_NAME + fetch_cluster_info_mock.assert_called_once() @mock.patch.dict(os.environ, {}) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") + @mock.patch(KUB_OPERATOR_EXEC) + @mock.patch(TEMP_FILE) + @mock.patch(f"{GKE_OP_PATH}.fetch_cluster_info") def test_execute_with_impersonation_service_chain_one_element( - self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock + self, fetch_cluster_info_mock, file_mock, exec_mock, get_con_mock ): - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) self.gke_op.impersonation_chain = ["test_account@example.com"] self.gke_op.execute(context=mock.MagicMock()) - mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() - - mock_execute_in_subprocess.assert_called_once_with( - [ - "gcloud", - "container", - "clusters", - "get-credentials", - CLUSTER_NAME, - "--project", - TEST_GCP_PROJECT_ID, - "--impersonate-service-account", - "test_account@example.com", - "--zone", - PROJECT_LOCATION, - ] - ) + fetch_cluster_info_mock.assert_called_once() + - assert self.gke_op.config_file == FILE_NAME +class TestGKEPodOperatorAsync: + def setup_method(self): + self.gke_op = GKEStartPodOperator( + project_id=TEST_GCP_PROJECT_ID, + location=PROJECT_LOCATION, + cluster_name=CLUSTER_NAME, + task_id=PROJECT_TASK_ID, + name=TASK_NAME, + namespace=NAMESPACE, + image=IMAGE, + deferrable=True, + ) + self.gke_op.pod = mock.MagicMock( + name=TASK_NAME, + namespace=NAMESPACE, + ) + self.gke_op._cluster_url = CLUSTER_URL + self.gke_op._ssl_ca_cert = SSL_CA_CERT + @mock.patch(KUB_OP_PATH.format("convert_config_file_to_dict")) @mock.patch.dict(os.environ, {}) + @mock.patch(KUB_OP_PATH.format("build_pod_request_obj")) + @mock.patch(KUB_OP_PATH.format("get_or_create_pod")) @mock.patch( "airflow.hooks.base.BaseHook.get_connections", return_value=[Connection(extra=json.dumps({"keyfile_dict": '{"private_key": "r4nd0m_k3y"}'}))], ) - @mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook") - @mock.patch("airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess") - @mock.patch("tempfile.NamedTemporaryFile") - def test_execute_with_impersonation_service_chain_more_elements( - self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock + @mock.patch(f"{GKE_OP_PATH}.fetch_cluster_info") + def test_async_create_pod_should_execute_successfully( + self, fetch_cluster_info_mock, get_con_mock, mocked_pod, mocked_pod_obj, mocked_config ): - type(file_mock.return_value.__enter__.return_value).name = PropertyMock( - side_effect=[FILE_NAME, "/path/to/new-file"] - ) - self.gke_op.impersonation_chain = ["test_account@example.com", "test_account1@example.com"] - with pytest.raises( - AirflowException, - match="Chained list of accounts is not supported, please specify only one service account", - ): - self.gke_op.execute(None) + """ + Asserts that a task is deferred and the GKEPodTrigger will be fired + when the GKEStartPodOperator is executed in deferrable mode when deferrable=True. + """ + with pytest.raises(TaskDeferred) as exc: + self.gke_op._cluster_url = CLUSTER_URL + self.gke_op._ssl_ca_cert = SSL_CA_CERT + self.gke_op.execute(context=mock.MagicMock()) + fetch_cluster_info_mock.assert_called_once() + assert isinstance(exc.value.trigger, GKEPodTrigger) diff --git a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py index 1bb5926b2c58c..459f9c2e5fbfd 100644 --- a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py @@ -20,12 +20,16 @@ import asyncio import logging import sys -from asyncio import Future +from asyncio import CancelledError, Future +from datetime import datetime import pytest +import pytz +from kubernetes.client import models as k8s from google.cloud.container_v1.types import Operation -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger +from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ContainerState +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEPodTrigger, GKEOperationTrigger from airflow.triggers.base import TriggerEvent if sys.version_info < (3, 8): @@ -33,6 +37,22 @@ else: from unittest import mock +TRIGGER_GKE_PATH = "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEPodTrigger" +TRIGGER_KUB_PATH = "airflow.providers.cncf.kubernetes.triggers.kubernetes_pod.KubernetesPodTrigger" +HOOK_PATH = "airflow.providers.google.cloud.hooks.kubernetes_engine.AsyncGKEPodHook" +POD_NAME = "test-pod-name" +NAMESPACE = "default" +POLL_INTERVAL = 2 +CLUSTER_CONTEXT = "test-context" +IN_CLUSTER = False +SHOULD_DELETE_POD = True +GET_LOGS = True +STARTUP_TIMEOUT_SECS = 120 +TRIGGER_START_TIME = datetime.now(tz=pytz.UTC) +CLUSTER_URL = "https://test-host" +SSL_CA_CERT = "TEST_SSL_CA_CERT_CONTENT" +FAILED_RESULT_MSG = "Test message that appears when trigger have failed event." +BASE_CONTAINER_NAME = "base" OPERATION_NAME = "test-operation-name" PROJECT_ID = "test-project-id" @@ -40,11 +60,241 @@ GCP_CONN_ID = "test-non-existing-project-id" DELEGATE_TO = "test-delegate-to" IMPERSONATION_CHAIN = ["impersonate", "this", "test"] -POLL_INTERVAL = 5 TRIGGER_PATH = "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEOperationTrigger" EXC_MSG = "test error msg" +@pytest.fixture +def trigger(): + return GKEPodTrigger( + pod_name=POD_NAME, + pod_namespace=NAMESPACE, + poll_interval=POLL_INTERVAL, + cluster_context=CLUSTER_CONTEXT, + in_cluster=IN_CLUSTER, + should_delete_pod=SHOULD_DELETE_POD, + get_logs=GET_LOGS, + startup_timeout=STARTUP_TIMEOUT_SECS, + trigger_start_time=TRIGGER_START_TIME, + cluster_url=CLUSTER_URL, + ssl_ca_cert=SSL_CA_CERT, + base_container_name=BASE_CONTAINER_NAME, + ) + + +class TestGKEPodTrigger: + @staticmethod + def _mock_pod_result(result_to_mock): + f = Future() + f.set_result(result_to_mock) + return f + + def test_serialize_should_execute_successfully(self, trigger): + classpath, kwargs_dict = trigger.serialize() + + assert classpath == TRIGGER_GKE_PATH + assert kwargs_dict == { + "pod_name": POD_NAME, + "pod_namespace": NAMESPACE, + "poll_interval": POLL_INTERVAL, + "cluster_context": CLUSTER_CONTEXT, + "in_cluster": IN_CLUSTER, + "should_delete_pod": SHOULD_DELETE_POD, + "get_logs": GET_LOGS, + "startup_timeout": STARTUP_TIMEOUT_SECS, + "trigger_start_time": TRIGGER_START_TIME, + "cluster_url": CLUSTER_URL, + "ssl_ca_cert": SSL_CA_CERT, + "base_container_name": BASE_CONTAINER_NAME, + } + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state") + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_run_loop_return_success_event_should_execute_successfully( + self, mock_hook, mock_method, trigger + ): + mock_hook.return_value.get_pod.return_value = self._mock_pod_result(mock.MagicMock()) + mock_method.return_value = ContainerState.TERMINATED + + expected_event = TriggerEvent( + { + "name": POD_NAME, + "namespace": NAMESPACE, + "status": "success", + "message": "All containers inside pod have started successfully.", + } + ) + actual_event = await (trigger.run()).asend(None) + + assert actual_event == expected_event + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state") + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_run_loop_return_failed_event_should_execute_successfully( + self, mock_hook, mock_method, trigger + ): + mock_hook.return_value.get_pod.return_value = self._mock_pod_result( + mock.MagicMock( + status=mock.MagicMock( + message=FAILED_RESULT_MSG, + ) + ) + ) + mock_method.return_value = ContainerState.FAILED + + expected_event = TriggerEvent( + { + "name": POD_NAME, + "namespace": NAMESPACE, + "status": "failed", + "message": FAILED_RESULT_MSG, + } + ) + actual_event = await (trigger.run()).asend(None) + + assert actual_event == expected_event + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state") + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_run_loop_return_waiting_event_should_execute_successfully( + self, mock_hook, mock_method, trigger, caplog + ): + mock_hook.return_value.get_pod.return_value = self._mock_pod_result(mock.MagicMock()) + mock_method.return_value = ContainerState.WAITING + + caplog.set_level(logging.INFO) + + task = asyncio.create_task(trigger.run().__anext__()) + await asyncio.sleep(0.5) + + assert not task.done() + assert "Container is not completed and still working." + assert f"Sleeping for {POLL_INTERVAL} seconds." + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state") + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_run_loop_return_running_event_should_execute_successfully( + self, mock_hook, mock_method, trigger, caplog + ): + mock_hook.return_value.get_pod.return_value = self._mock_pod_result(mock.MagicMock()) + mock_method.return_value = ContainerState.RUNNING + + caplog.set_level(logging.INFO) + + task = asyncio.create_task(trigger.run().__anext__()) + await asyncio.sleep(0.5) + + assert not task.done() + assert "Container is not completed and still working." + assert f"Sleeping for {POLL_INTERVAL} seconds." + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_logging_in_trigger_when_exception_should_execute_successfully( + self, mock_hook, trigger, caplog + ): + """ + Test that GKEPodTrigger fires the correct event in case of an error. + """ + mock_hook.return_value.get_pod.side_effect = Exception("Test exception") + + generator = trigger.run() + actual = await generator.asend(None) + assert ( + TriggerEvent( + {"name": POD_NAME, "namespace": NAMESPACE, "status": "error", "message": "Test exception"} + ) + == actual + ) + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_KUB_PATH}.define_container_state") + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_logging_in_trigger_when_fail_should_execute_successfully( + self, mock_hook, mock_method, trigger, caplog + ): + """ + Test that GKEPodTrigger fires the correct event in case of fail. + """ + mock_hook.return_value.get_pod.return_value = self._mock_pod_result(mock.MagicMock()) + mock_method.return_value = ContainerState.FAILED + caplog.set_level(logging.INFO) + + generator = trigger.run() + await generator.asend(None) + assert "Container logs:" + + @pytest.mark.asyncio + @mock.patch(f"{TRIGGER_GKE_PATH}._get_async_hook") + async def test_logging_in_trigger_when_cancelled_should_execute_successfully( + self, mock_hook, trigger, caplog + ): + """ + Test that GKEPodTrigger fires the correct event in case if the task was cancelled. + """ + mock_hook.return_value.get_pod.side_effect = CancelledError() + mock_hook.return_value.read_logs.return_value = self._mock_pod_result(mock.MagicMock()) + mock_hook.return_value.delete_pod.return_value = self._mock_pod_result(mock.MagicMock()) + + generator = trigger.run() + actual = await generator.asend(None) + assert ( + TriggerEvent( + { + "name": POD_NAME, + "namespace": NAMESPACE, + "status": "cancelled", + "message": "Pod execution was cancelled", + } + ) + == actual + ) + assert "Outputting container logs..." in caplog.text + assert "Deleting pod..." in caplog.text + + @pytest.mark.parametrize( + "container_state, expected_state", + [ + ( + {"running": k8s.V1ContainerStateRunning(), "terminated": None, "waiting": None}, + ContainerState.RUNNING, + ), + ( + {"running": None, "terminated": k8s.V1ContainerStateTerminated(exit_code=0), "waiting": None}, + ContainerState.TERMINATED, + ), + ( + {"running": None, "terminated": None, "waiting": k8s.V1ContainerStateWaiting()}, + ContainerState.WAITING, + ), + ], + ) + def test_define_container_state_should_execute_successfully( + self, trigger, container_state, expected_state + ): + pod = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name="base", namespace="default"), + status=k8s.V1PodStatus( + container_statuses=[ + k8s.V1ContainerStatus( + name="base", + image="alpine", + image_id="1", + ready=True, + restart_count=1, + state=k8s.V1ContainerState(**container_state), + ) + ] + ), + ) + + assert expected_state == trigger.define_container_state(pod) + + @pytest.fixture def trigger(): return GKEOperationTrigger( diff --git a/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py b/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py index 2f4dadfd952ed..8d031c5d917d6 100644 --- a/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py +++ b/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py @@ -68,9 +68,11 @@ in_cluster=False, is_delete_operator_pod=True, get_logs=True, + deferrable=True, ) - pod_task_xcom = GKEStartPodOperator( + # [START howto_operator_gke_start_pod_xcom_async] + pod_task_xcom_async = GKEStartPodOperator( task_id="pod_task_xcom", project_id=GCP_PROJECT_ID, location=GCP_LOCATION, @@ -82,13 +84,17 @@ in_cluster=False, is_delete_operator_pod=True, do_xcom_push=True, + deferrable=True, get_logs=True, ) + # [END howto_operator_gke_start_pod_xcom_async] + # [START howto_operator_gke_xcom_result_async] pod_task_xcom_result = BashOperator( bash_command="echo \"{{ task_instance.xcom_pull('pod_task_xcom')[0] }}\"", task_id="pod_task_xcom_result", ) + # [END howto_operator_gke_xcom_result_async] # [START howto_operator_gke_delete_cluster_async] delete_cluster = GKEDeleteClusterOperator( @@ -101,8 +107,8 @@ # [END howto_operator_gke_delete_cluster_async] create_cluster >> pod_task >> delete_cluster - create_cluster >> pod_task_xcom >> delete_cluster - pod_task_xcom >> pod_task_xcom_result + create_cluster >> pod_task_xcom_async >> delete_cluster + pod_task_xcom_async >> pod_task_xcom_result from tests.system.utils.watcher import watcher From 4e53fbabbca9a62dc060f9e54f8d3b435d8d9c81 Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Wed, 1 Feb 2023 13:29:20 +0100 Subject: [PATCH 2/5] Change naming for GKEHook and add comments --- RELEASE_NOTES.rst | 6 +++--- airflow/contrib/hooks/__init__.py | 1 - airflow/providers/google/cloud/hooks/kubernetes_engine.py | 6 +++--- .../google/cloud/operators/test_kubernetes_engine.py | 2 +- 4 files changed, 7 insertions(+), 8 deletions(-) diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index 6da321c799cc2..54e68d691d8f0 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -4765,7 +4765,7 @@ The following table shows changes in import paths. - ``airflow.providers.google.cloud.hooks.bigtable.BigtableHook`` * - ``airflow.contrib.hooks.gcp_cloud_build_hook.CloudBuildHook`` - ``airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook`` - * - ``airflow.contrib.hooks.gcp_container_hook.GKEClusterHook`` + * - ``airflow.contrib.hooks.gcp_container_hook.GKEHook`` - ``airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook`` * - ``airflow.contrib.hooks.gcp_compute_hook.GceHook`` - ``airflow.providers.google.cloud.hooks.compute.ComputeEngineHook`` @@ -7841,7 +7841,7 @@ Bug fixes - [AIRFLOW-4524] Fix bug with "Ignore \*" toggles in RBAC mode (#5378) - [AIRFLOW-4765] Fix DataProcPigOperator execute method (#5426) - [AIRFLOW-4798] Obviate interdependencies for dagbag and TI tests (#5422) -- [AIRFLOW-4800] Fix GKEClusterHook ctor calls (#5424) +- [AIRFLOW-4800] Fix GKEHook ctor calls (#5424) - [AIRFLOW-4799] Don't mutate self.env in BashOperator execute method (#5421) - [AIRFLOW-4393] Add retry logic when fetching pod status and/or logs in KubernetesPodOperator (#5284) - [AIRFLOW-4174] Fix HttpHook run with backoff (#5213) @@ -8788,7 +8788,7 @@ Bug fixes - [AIRFLOW-2645][AIRFLOW-2617] Add worker_container_image_pull_policy - [AIRFLOW-2661] fix config dags_volume_subpath and logs_volume_subpath - [AIRFLOW-3550] Standardize GKE hook (#4364) -- [AIRFLOW-2863] Fix GKEClusterHook catching wrong exception (#3711) +- [AIRFLOW-2863] Fix GKEHook catching wrong exception (#3711) - [AIRFLOW-2939][AIRFLOW-3568] Fix TypeError in GCSToS3Op & S3ToGCSOp (#4371) - [AIRFLOW-3327] Add support for location in BigQueryHook (#4324) - [AIRFLOW-3438] Fix default values in BigQuery Hook & BigQueryOperator (… diff --git a/airflow/contrib/hooks/__init__.py b/airflow/contrib/hooks/__init__.py index cffc937097306..7a46fc09ab7c1 100644 --- a/airflow/contrib/hooks/__init__.py +++ b/airflow/contrib/hooks/__init__.py @@ -154,7 +154,6 @@ }, 'gcp_container_hook': { 'GKEHook': 'airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook', - 'GKEClusterHook': 'airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook', }, 'gcp_dataflow_hook': { 'DataflowHook': 'airflow.providers.google.cloud.hooks.dataflow.DataflowHook', diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index c5645623ef497..62699ede77aaf 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -64,7 +64,7 @@ class GKEHook(GoogleBaseHook): """ - Hook for Google Kubernetes Engine APIs. + Hook for managing Google Kubernetes Engine cluster APIs. All the methods in the hook where project_id is used must be called with keyword arguments rather than positional. @@ -346,7 +346,7 @@ async def get_operation( class GKEPodHook(GoogleBaseHook): - """Hook for managing pods.""" + """Hook for managing Google Kubernetes Engine pod APIs.""" def __init__( self, @@ -416,7 +416,7 @@ def get_pod(self, name: str, namespace: str) -> V1Pod: class AsyncGKEPodHook(GoogleBaseAsyncHook): """ - Hook for managing GKE pods in asynchronous way. + Hook for managing Google Kubernetes Engine pods APIs in asynchronous way. :param cluster_url: The URL pointed to the cluster. :param ssl_ca_cert: SSL certificate that is used for authentication to the pod. diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index abc006b6d9d86..5ee7f94743b32 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -61,7 +61,7 @@ FILE_NAME = "/tmp/mock_name" KUB_OP_PATH = "airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.{}" GKE_HOOK_MODULE_PATH = "airflow.providers.google.cloud.operators.kubernetes_engine" -GKE_HOOK_PATH = f"{GKE_HOOK_MODULE_PATH}.GKEClusterHook" +GKE_HOOK_PATH = f"{GKE_HOOK_MODULE_PATH}.GKEHook" KUB_OPERATOR_EXEC = "airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute" TEMP_FILE = "tempfile.NamedTemporaryFile" GKE_OP_PATH = "airflow.providers.google.cloud.operators.kubernetes_engine.GKEStartPodOperator" From 61aaae141bce01258d978141bdc7295eb787fa1f Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Mon, 27 Feb 2023 08:10:01 +0100 Subject: [PATCH 3/5] Rebase main, revert unrelated changes --- RELEASE_NOTES.rst | 6 +-- airflow/contrib/hooks/__init__.py | 1 + .../google/cloud/hooks/kubernetes_engine.py | 2 +- .../cloud/triggers/kubernetes_engine.py | 3 +- .../cloud/hooks/test_kubernetes_engine.py | 2 +- .../cloud/triggers/test_kubernetes_engine.py | 38 ++++++++++--------- 6 files changed, 28 insertions(+), 24 deletions(-) diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index 54e68d691d8f0..6da321c799cc2 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -4765,7 +4765,7 @@ The following table shows changes in import paths. - ``airflow.providers.google.cloud.hooks.bigtable.BigtableHook`` * - ``airflow.contrib.hooks.gcp_cloud_build_hook.CloudBuildHook`` - ``airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook`` - * - ``airflow.contrib.hooks.gcp_container_hook.GKEHook`` + * - ``airflow.contrib.hooks.gcp_container_hook.GKEClusterHook`` - ``airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook`` * - ``airflow.contrib.hooks.gcp_compute_hook.GceHook`` - ``airflow.providers.google.cloud.hooks.compute.ComputeEngineHook`` @@ -7841,7 +7841,7 @@ Bug fixes - [AIRFLOW-4524] Fix bug with "Ignore \*" toggles in RBAC mode (#5378) - [AIRFLOW-4765] Fix DataProcPigOperator execute method (#5426) - [AIRFLOW-4798] Obviate interdependencies for dagbag and TI tests (#5422) -- [AIRFLOW-4800] Fix GKEHook ctor calls (#5424) +- [AIRFLOW-4800] Fix GKEClusterHook ctor calls (#5424) - [AIRFLOW-4799] Don't mutate self.env in BashOperator execute method (#5421) - [AIRFLOW-4393] Add retry logic when fetching pod status and/or logs in KubernetesPodOperator (#5284) - [AIRFLOW-4174] Fix HttpHook run with backoff (#5213) @@ -8788,7 +8788,7 @@ Bug fixes - [AIRFLOW-2645][AIRFLOW-2617] Add worker_container_image_pull_policy - [AIRFLOW-2661] fix config dags_volume_subpath and logs_volume_subpath - [AIRFLOW-3550] Standardize GKE hook (#4364) -- [AIRFLOW-2863] Fix GKEHook catching wrong exception (#3711) +- [AIRFLOW-2863] Fix GKEClusterHook catching wrong exception (#3711) - [AIRFLOW-2939][AIRFLOW-3568] Fix TypeError in GCSToS3Op & S3ToGCSOp (#4371) - [AIRFLOW-3327] Add support for location in BigQueryHook (#4324) - [AIRFLOW-3438] Fix default values in BigQuery Hook & BigQueryOperator (… diff --git a/airflow/contrib/hooks/__init__.py b/airflow/contrib/hooks/__init__.py index 7a46fc09ab7c1..cffc937097306 100644 --- a/airflow/contrib/hooks/__init__.py +++ b/airflow/contrib/hooks/__init__.py @@ -154,6 +154,7 @@ }, 'gcp_container_hook': { 'GKEHook': 'airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook', + 'GKEClusterHook': 'airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook', }, 'gcp_dataflow_hook': { 'DataflowHook': 'airflow.providers.google.cloud.hooks.dataflow.DataflowHook', diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 62699ede77aaf..11d6bce45a36e 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -33,7 +33,7 @@ import google.auth.credentials from gcloud.aio.auth import Token -from google.api_core.exceptions import AlreadyExists, NotFound +from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry from google.auth.transport import requests as google_requests diff --git a/airflow/providers/google/cloud/triggers/kubernetes_engine.py b/airflow/providers/google/cloud/triggers/kubernetes_engine.py index 5508abc9975fc..673fa061676f9 100644 --- a/airflow/providers/google/cloud/triggers/kubernetes_engine.py +++ b/airflow/providers/google/cloud/triggers/kubernetes_engine.py @@ -28,8 +28,8 @@ except ImportError: # preserve backward compatibility for older versions of cncf.kubernetes provider from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import KubernetesPodTrigger +from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, AsyncGKEPodHook from airflow.triggers.base import BaseTrigger, TriggerEvent -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEPodHook, AsyncGKEHook class GKEPodTrigger(KubernetesPodTrigger): @@ -213,4 +213,3 @@ def _get_hook(self) -> AsyncGKEHook: impersonation_chain=self.impersonation_chain, ) return self._hook - diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index da1a31aaea9ad..64807649734df 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -26,7 +26,7 @@ from google.cloud.container_v1.types import Cluster from airflow.exceptions import AirflowException -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, GKEHook, AsyncGKEPodHook +from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, AsyncGKEPodHook, GKEHook from airflow.providers.google.common.consts import CLIENT_INFO from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id diff --git a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py index 459f9c2e5fbfd..9013980ecfac6 100644 --- a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py @@ -25,11 +25,11 @@ import pytest import pytz -from kubernetes.client import models as k8s from google.cloud.container_v1.types import Operation +from kubernetes.client import models as k8s from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ContainerState -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEPodTrigger, GKEOperationTrigger +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEPodTrigger from airflow.triggers.base import TriggerEvent if sys.version_info < (3, 8): @@ -296,7 +296,7 @@ def test_define_container_state_should_execute_successfully( @pytest.fixture -def trigger(): +def operation_trigger(): return GKEOperationTrigger( operation_name=OPERATION_NAME, project_id=PROJECT_ID, @@ -321,8 +321,8 @@ def func(**kwargs): class TestGKEOperationTrigger: - def test_serialize(self, trigger): - classpath, trigger_init_kwargs = trigger.serialize() + def test_serialize(self, operation_trigger): + classpath, trigger_init_kwargs = operation_trigger.serialize() assert classpath == TRIGGER_PATH assert trigger_init_kwargs == { "operation_name": OPERATION_NAME, @@ -336,7 +336,9 @@ def test_serialize(self, trigger): @pytest.mark.asyncio @mock.patch(f"{TRIGGER_PATH}._get_hook") - async def test_run_loop_return_success_event(self, mock_hook, trigger, async_get_operation_result): + async def test_run_loop_return_success_event( + self, mock_hook, operation_trigger, async_get_operation_result + ): mock_hook.return_value.get_operation.return_value = async_get_operation_result( name=OPERATION_NAME, status=Operation.Status.DONE, @@ -349,7 +351,7 @@ async def test_run_loop_return_success_event(self, mock_hook, trigger, async_get "operation_name": OPERATION_NAME, } ) - actual_event = await (trigger.run()).asend(None) + actual_event = await (operation_trigger.run()).asend(None) assert actual_event == expected_event @@ -358,7 +360,7 @@ async def test_run_loop_return_success_event(self, mock_hook, trigger, async_get async def test_run_loop_return_failed_event_status_unspecified( self, mock_hook, - trigger, + operation_trigger, async_get_operation_result, ): mock_hook.return_value.get_operation.return_value = async_get_operation_result( @@ -372,7 +374,7 @@ async def test_run_loop_return_failed_event_status_unspecified( "message": f"Operation has failed with status: {Operation.Status.STATUS_UNSPECIFIED}", } ) - actual_event = await (trigger.run()).asend(None) + actual_event = await (operation_trigger.run()).asend(None) assert actual_event == expected_event @@ -381,7 +383,7 @@ async def test_run_loop_return_failed_event_status_unspecified( async def test_run_loop_return_failed_event_status_aborting( self, mock_hook, - trigger, + operation_trigger, async_get_operation_result, ): mock_hook.return_value.get_operation.return_value = async_get_operation_result( @@ -395,13 +397,15 @@ async def test_run_loop_return_failed_event_status_aborting( "message": f"Operation has failed with status: {Operation.Status.ABORTING}", } ) - actual_event = await (trigger.run()).asend(None) + actual_event = await (operation_trigger.run()).asend(None) assert actual_event == expected_event @pytest.mark.asyncio @mock.patch(f"{TRIGGER_PATH}._get_hook") - async def test_run_loop_return_error_event(self, mock_hook, trigger, async_get_operation_result): + async def test_run_loop_return_error_event( + self, mock_hook, operation_trigger, async_get_operation_result + ): mock_hook.return_value.get_operation.side_effect = Exception(EXC_MSG) expected_event = TriggerEvent( @@ -410,7 +414,7 @@ async def test_run_loop_return_error_event(self, mock_hook, trigger, async_get_o "message": EXC_MSG, } ) - actual_event = await (trigger.run()).asend(None) + actual_event = await (operation_trigger.run()).asend(None) assert actual_event == expected_event @@ -419,7 +423,7 @@ async def test_run_loop_return_error_event(self, mock_hook, trigger, async_get_o async def test_run_loop_return_waiting_event_pending_status( self, mock_hook, - trigger, + operation_trigger, async_get_operation_result, caplog, ): @@ -430,7 +434,7 @@ async def test_run_loop_return_waiting_event_pending_status( caplog.set_level(logging.INFO) - task = asyncio.create_task(trigger.run().__anext__()) + task = asyncio.create_task(operation_trigger.run().__anext__()) await asyncio.sleep(0.5) assert not task.done() @@ -442,7 +446,7 @@ async def test_run_loop_return_waiting_event_pending_status( async def test_run_loop_return_waiting_event_running_status( self, mock_hook, - trigger, + operation_trigger, async_get_operation_result, caplog, ): @@ -453,7 +457,7 @@ async def test_run_loop_return_waiting_event_running_status( caplog.set_level(logging.INFO) - task = asyncio.create_task(trigger.run().__anext__()) + task = asyncio.create_task(operation_trigger.run().__anext__()) await asyncio.sleep(0.5) assert not task.done() From cf70dc649d72a0190408d2234c53cf399d8f2b2b Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Tue, 21 Mar 2023 10:54:57 +0100 Subject: [PATCH 4/5] Add review suggestions + rebase --- .../google/cloud/hooks/kubernetes_engine.py | 4 ++-- .../cloud/operators/kubernetes_engine.py | 4 ++-- .../google/cloud/triggers/kubernetes_engine.py | 16 ++++++++-------- .../operators/cloud/kubernetes_engine.rst | 15 ++++++++++++--- .../cloud/hooks/test_kubernetes_engine.py | 18 +++++++++--------- .../cloud/operators/test_kubernetes_engine.py | 6 +++--- .../cloud/triggers/test_kubernetes_engine.py | 16 ++++++++-------- .../example_kubernetes_engine_async.py | 2 +- 8 files changed, 45 insertions(+), 36 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 11d6bce45a36e..704988f52d3a3 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -295,7 +295,7 @@ def get_cluster( ) -class AsyncGKEHook(GoogleBaseAsyncHook): +class GKEAsyncHook(GoogleBaseAsyncHook): """Hook implemented with usage of asynchronous client of GKE.""" sync_hook_class = GKEHook @@ -414,7 +414,7 @@ def get_pod(self, name: str, namespace: str) -> V1Pod: ) -class AsyncGKEPodHook(GoogleBaseAsyncHook): +class GKEPodAsyncHook(GoogleBaseAsyncHook): """ Hook for managing Google Kubernetes Engine pods APIs in asynchronous way. diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index 84466f27adbe8..f5f5e55b89aa6 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -39,7 +39,7 @@ KubernetesEnginePodLink, ) from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEPodTrigger +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEStartPodTrigger from airflow.utils.timezone import utcnow if TYPE_CHECKING: @@ -515,7 +515,7 @@ def invoke_defer_method(self): """Method to easily redefine triggers which are being used in child classes.""" trigger_start_time = utcnow() self.defer( - trigger=GKEPodTrigger( + trigger=GKEStartPodTrigger( pod_name=self.pod.metadata.name, pod_namespace=self.pod.metadata.namespace, trigger_start_time=trigger_start_time, diff --git a/airflow/providers/google/cloud/triggers/kubernetes_engine.py b/airflow/providers/google/cloud/triggers/kubernetes_engine.py index 673fa061676f9..0e8baec6f9f7a 100644 --- a/airflow/providers/google/cloud/triggers/kubernetes_engine.py +++ b/airflow/providers/google/cloud/triggers/kubernetes_engine.py @@ -28,11 +28,11 @@ except ImportError: # preserve backward compatibility for older versions of cncf.kubernetes provider from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import KubernetesPodTrigger -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, AsyncGKEPodHook +from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEAsyncHook, GKEPodAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent -class GKEPodTrigger(KubernetesPodTrigger): +class GKEStartPodTrigger(KubernetesPodTrigger): """ Trigger for checking pod status until it finishes its job. @@ -96,7 +96,7 @@ def __init__( def serialize(self) -> tuple[str, dict[str, Any]]: return ( - "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEPodTrigger", + "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEStartPodTrigger", { "pod_name": self.pod_name, "pod_namespace": self.pod_namespace, @@ -113,8 +113,8 @@ def serialize(self) -> tuple[str, dict[str, Any]]: }, ) - def _get_async_hook(self) -> AsyncGKEPodHook: # type: ignore[override] - return AsyncGKEPodHook( + def _get_async_hook(self) -> GKEPodAsyncHook: # type: ignore[override] + return GKEPodAsyncHook( cluster_url=self._cluster_url, ssl_ca_cert=self._ssl_ca_cert, ) @@ -143,7 +143,7 @@ def __init__( self.impersonation_chain = impersonation_chain self.poll_interval = poll_interval - self._hook: AsyncGKEHook | None = None + self._hook: GKEAsyncHook | None = None def serialize(self) -> tuple[str, dict[str, Any]]: """Serializes GKEOperationTrigger arguments and classpath.""" @@ -204,9 +204,9 @@ async def run(self) -> AsyncIterator["TriggerEvent"]: # type: ignore[override] ) return - def _get_hook(self) -> AsyncGKEHook: + def _get_hook(self) -> GKEAsyncHook: if self._hook is None: - self._hook = AsyncGKEHook( + self._hook = GKEAsyncHook( gcp_conn_id=self.gcp_conn_id, location=self.location, delegate_to=self.delegate_to, diff --git a/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst b/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst index b0545a5b3d7d1..013cb15328f0a 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst @@ -59,7 +59,10 @@ definition, is required when creating a cluster with :start-after: [START howto_operator_gke_create_cluster] :end-before: [END howto_operator_gke_create_cluster] -You can use deferrable mode for this action in order to run the operator asynchronously: +You can use deferrable mode for this action in order to run the operator asynchronously. It will give you a +possibility to free up the worker when it knows it has to wait, and hand off the job of resuming Operator to a Trigger. +As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a +lot less resources wasted on idle Operators or Sensors: .. exampleinclude:: /../../tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py :language: python @@ -83,7 +86,10 @@ This would also delete all the nodes allocated to the cluster. :start-after: [START howto_operator_gke_delete_cluster] :end-before: [END howto_operator_gke_delete_cluster] -You can use deferrable mode for this action in order to run the operator asynchronously: +You can use deferrable mode for this action in order to run the operator asynchronously. It will give you a +possibility to free up the worker when it knows it has to wait, and hand off the job of resuming Operator to a Trigger. +As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a +lot less resources wasted on idle Operators or Sensors: .. exampleinclude:: /../../tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py :language: python @@ -148,7 +154,10 @@ And then use it in other operators: :start-after: [START howto_operator_gke_xcom_result] :end-before: [END howto_operator_gke_xcom_result] -You can use deferrable mode for this action in order to run the operator asynchronously: +You can use deferrable mode for this action in order to run the operator asynchronously. It will give you a +possibility to free up the worker when it knows it has to wait, and hand off the job of resuming Operator to a Trigger. +As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a +lot less resources wasted on idle Operators or Sensors: .. exampleinclude:: /../../tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py :language: python diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index 64807649734df..6aa468bc2aa3e 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -26,7 +26,7 @@ from google.cloud.container_v1.types import Cluster from airflow.exceptions import AirflowException -from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook, AsyncGKEPodHook, GKEHook +from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEAsyncHook, GKEHook, GKEPodAsyncHook from airflow.providers.google.common.consts import CLIENT_INFO from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id @@ -45,7 +45,7 @@ SSL_CA_CERT = "test-ssl-ca-cert" POD_NAME = "test-pod-name" POD_NAMESPACE = "test" -ASYNC_HOOK_STRING = GKE_STRING.format("AsyncGKEHook") +ASYNC_HOOK_STRING = GKE_STRING.format("GKEAsyncHook") GCP_CONN_ID = "test-gcp-conn-id" DELEGATE_TO = "test-delegate-to" IMPERSONATE_CHAIN = ["impersonate", "this", "test"] @@ -296,7 +296,7 @@ def test_wait_for_response_running(self, time_mock, operation_mock): assert operation_mock.call_count == 2 -class TestAsyncGKEPodHook: +class TestGKEPodAsyncHook: @staticmethod def make_mock_awaitable(mock_obj, result=None): f = Future() @@ -306,14 +306,14 @@ def make_mock_awaitable(mock_obj, result=None): @pytest.fixture() def async_hook(self): - return AsyncGKEPodHook( + return GKEPodAsyncHook( cluster_url=CLUSTER_URL, ssl_ca_cert=SSL_CA_CERT, ) @pytest.mark.asyncio @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) - @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod")) async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook): self.make_mock_awaitable(read_namespace_pod_mock) @@ -328,7 +328,7 @@ async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook) @pytest.mark.asyncio @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) - @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.delete_namespaced_pod")) async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook): self.make_mock_awaitable(delete_namespaced_pod) @@ -344,7 +344,7 @@ async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook @pytest.mark.asyncio @mock.patch(GKE_STRING.format("Token"), mock.MagicMock()) - @mock.patch(GKE_STRING.format("AsyncGKEPodHook.get_conn")) + @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn")) @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod_log")) async def test_read_logs(self, read_namespaced_pod_log, get_conn_mock, async_hook, caplog): self.make_mock_awaitable(read_namespaced_pod_log, result="Test string #1\nTest string #2\n") @@ -364,7 +364,7 @@ async def test_read_logs(self, read_namespaced_pod_log, get_conn_mock, async_hoo @pytest.fixture() def async_gke_hook(): - return AsyncGKEHook( + return GKEAsyncHook( gcp_conn_id=GCP_CONN_ID, delegate_to=DELEGATE_TO, location=GKE_ZONE, @@ -381,7 +381,7 @@ def mock_async_gke_cluster_client(): return client -class TestAsyncGKEHook: +class TestGKEAsyncHook: @staticmethod def make_get_client_awaitable(mock_obj, result): if sys.version_info < (3, 8): diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index 5ee7f94743b32..a474c8d47455a 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -31,7 +31,7 @@ GKEDeleteClusterOperator, GKEStartPodOperator, ) -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEPodTrigger +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEStartPodTrigger TEST_GCP_PROJECT_ID = "test-id" PROJECT_LOCATION = "test-location" @@ -326,7 +326,7 @@ def test_async_create_pod_should_execute_successfully( self, fetch_cluster_info_mock, get_con_mock, mocked_pod, mocked_pod_obj, mocked_config ): """ - Asserts that a task is deferred and the GKEPodTrigger will be fired + Asserts that a task is deferred and the GKEStartPodTrigger will be fired when the GKEStartPodOperator is executed in deferrable mode when deferrable=True. """ with pytest.raises(TaskDeferred) as exc: @@ -334,4 +334,4 @@ def test_async_create_pod_should_execute_successfully( self.gke_op._ssl_ca_cert = SSL_CA_CERT self.gke_op.execute(context=mock.MagicMock()) fetch_cluster_info_mock.assert_called_once() - assert isinstance(exc.value.trigger, GKEPodTrigger) + assert isinstance(exc.value.trigger, GKEStartPodTrigger) diff --git a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py index 9013980ecfac6..eb3497012f97b 100644 --- a/tests/providers/google/cloud/triggers/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/triggers/test_kubernetes_engine.py @@ -29,7 +29,7 @@ from kubernetes.client import models as k8s from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ContainerState -from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEPodTrigger +from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEStartPodTrigger from airflow.triggers.base import TriggerEvent if sys.version_info < (3, 8): @@ -37,9 +37,9 @@ else: from unittest import mock -TRIGGER_GKE_PATH = "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEPodTrigger" +TRIGGER_GKE_PATH = "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEStartPodTrigger" TRIGGER_KUB_PATH = "airflow.providers.cncf.kubernetes.triggers.kubernetes_pod.KubernetesPodTrigger" -HOOK_PATH = "airflow.providers.google.cloud.hooks.kubernetes_engine.AsyncGKEPodHook" +HOOK_PATH = "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEPodAsyncHook" POD_NAME = "test-pod-name" NAMESPACE = "default" POLL_INTERVAL = 2 @@ -66,7 +66,7 @@ @pytest.fixture def trigger(): - return GKEPodTrigger( + return GKEStartPodTrigger( pod_name=POD_NAME, pod_namespace=NAMESPACE, poll_interval=POLL_INTERVAL, @@ -82,7 +82,7 @@ def trigger(): ) -class TestGKEPodTrigger: +class TestGKEStartPodTrigger: @staticmethod def _mock_pod_result(result_to_mock): f = Future() @@ -198,7 +198,7 @@ async def test_logging_in_trigger_when_exception_should_execute_successfully( self, mock_hook, trigger, caplog ): """ - Test that GKEPodTrigger fires the correct event in case of an error. + Test that GKEStartPodTrigger fires the correct event in case of an error. """ mock_hook.return_value.get_pod.side_effect = Exception("Test exception") @@ -218,7 +218,7 @@ async def test_logging_in_trigger_when_fail_should_execute_successfully( self, mock_hook, mock_method, trigger, caplog ): """ - Test that GKEPodTrigger fires the correct event in case of fail. + Test that GKEStartPodTrigger fires the correct event in case of fail. """ mock_hook.return_value.get_pod.return_value = self._mock_pod_result(mock.MagicMock()) mock_method.return_value = ContainerState.FAILED @@ -234,7 +234,7 @@ async def test_logging_in_trigger_when_cancelled_should_execute_successfully( self, mock_hook, trigger, caplog ): """ - Test that GKEPodTrigger fires the correct event in case if the task was cancelled. + Test that GKEStartPodTrigger fires the correct event in case if the task was cancelled. """ mock_hook.return_value.get_pod.side_effect = CancelledError() mock_hook.return_value.read_logs.return_value = self._mock_pod_result(mock.MagicMock()) diff --git a/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py b/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py index 8d031c5d917d6..13c310d880210 100644 --- a/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py +++ b/tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py @@ -73,7 +73,7 @@ # [START howto_operator_gke_start_pod_xcom_async] pod_task_xcom_async = GKEStartPodOperator( - task_id="pod_task_xcom", + task_id="pod_task_xcom_async", project_id=GCP_PROJECT_ID, location=GCP_LOCATION, cluster_name=CLUSTER_NAME, From deb9b891f40e2643070a7200acf95df5931a3c77 Mon Sep 17 00:00:00 2001 From: VladaZakharova Date: Tue, 4 Apr 2023 11:44:44 +0200 Subject: [PATCH 5/5] Add deprecation warning for deleted method + rebase --- .../google/cloud/operators/kubernetes_engine.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index f5f5e55b89aa6..21ef3692989e0 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -473,6 +473,15 @@ def __init__( if self.config_file: raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.") + @staticmethod + def get_gke_config_file(): + warnings.warn( + "The `get_gke_config_file` method is deprecated, " + "please use `fetch_cluster_info` instead to get the cluster info for connecting to it.", + DeprecationWarning, + stacklevel=1, + ) + @cached_property def cluster_hook(self) -> GKEHook: return GKEHook(