Source code for airflow.providers.google.cloud.hooks.cloud_composer

#
# 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

from typing import TYPE_CHECKING, Sequence

from google.api_core.client_options import ClientOptions
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
from google.cloud.orchestration.airflow.service_v1 import (
    EnvironmentsAsyncClient,
    EnvironmentsClient,
    ImageVersionsClient,
)

from airflow.exceptions import AirflowException
from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook

if TYPE_CHECKING:
    from google.api_core.operation import Operation
    from google.api_core.operation_async import AsyncOperation
    from google.api_core.retry import Retry
    from google.api_core.retry_async import AsyncRetry
    from google.cloud.orchestration.airflow.service_v1.services.environments.pagers import (
        ListEnvironmentsPager,
    )
    from google.cloud.orchestration.airflow.service_v1.services.image_versions.pagers import (
        ListImageVersionsPager,
    )
    from google.cloud.orchestration.airflow.service_v1.types import Environment
    from google.protobuf.field_mask_pb2 import FieldMask


[docs]class CloudComposerHook(GoogleBaseHook): """Hook for Google Cloud Composer APIs."""
[docs] client_options = ClientOptions(api_endpoint="composer.googleapis.com:443")
def __init__(self, **kwargs): if kwargs.get("delegate_to") is not None: raise RuntimeError( "The `delegate_to` parameter has been deprecated before and finally removed in this version" " of Google Provider. You MUST convert it to `impersonate_chain`" ) super().__init__(**kwargs)
[docs] def get_environment_client(self) -> EnvironmentsClient: """Retrieves client library object that allow access Environments service.""" return EnvironmentsClient( credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=self.client_options, )
[docs] def get_image_versions_client(self) -> ImageVersionsClient: """Retrieves client library object that allow access Image Versions service.""" return ImageVersionsClient( credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=self.client_options, )
[docs] def wait_for_operation(self, operation: Operation, timeout: float | None = None): """Waits for long-lasting operation to complete.""" try: return operation.result(timeout=timeout) except Exception: error = operation.exception(timeout=timeout) raise AirflowException(error)
[docs] def get_operation(self, operation_name): return self.get_environment_client().transport.operations_client.get_operation(name=operation_name)
[docs] def get_environment_name(self, project_id, region, environment_id): return f"projects/{project_id}/locations/{region}/environments/{environment_id}"
[docs] def get_parent(self, project_id, region): return f"projects/{project_id}/locations/{region}"
@GoogleBaseHook.fallback_to_default_project_id
[docs] def create_environment( self, project_id: str, region: str, environment: Environment | dict, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> Operation: """ Create a new environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment: The environment to create. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() result = client.create_environment( request={"parent": self.get_parent(project_id, region), "environment": environment}, retry=retry, timeout=timeout, metadata=metadata, ) return result
@GoogleBaseHook.fallback_to_default_project_id
[docs] def delete_environment( self, project_id: str, region: str, environment_id: str, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> Operation: """ Delete an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() name = self.get_environment_name(project_id, region, environment_id) result = client.delete_environment( request={"name": name}, retry=retry, timeout=timeout, metadata=metadata ) return result
@GoogleBaseHook.fallback_to_default_project_id
[docs] def get_environment( self, project_id: str, region: str, environment_id: str, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> Environment: """ Get an existing environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() result = client.get_environment( request={"name": self.get_environment_name(project_id, region, environment_id)}, retry=retry, timeout=timeout, metadata=metadata, ) return result
@GoogleBaseHook.fallback_to_default_project_id
[docs] def list_environments( self, project_id: str, region: str, page_size: int | None = None, page_token: str | None = None, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> ListEnvironmentsPager: """ List environments. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param page_size: The maximum number of environments to return. :param page_token: The next_page_token value returned from a previous List request, if any. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() result = client.list_environments( request={ "parent": self.get_parent(project_id, region), "page_size": page_size, "page_token": page_token, }, retry=retry, timeout=timeout, metadata=metadata, ) return result
@GoogleBaseHook.fallback_to_default_project_id
[docs] def update_environment( self, project_id: str, region: str, environment_id: str, environment: Environment | dict, update_mask: dict | FieldMask, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> Operation: r""" Update an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param environment: A patch environment. Fields specified by the ``updateMask`` will be copied from the patch environment into the environment under update. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param update_mask: Required. A comma-separated list of paths, relative to ``Environment``, of fields to update. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.protobuf.field_mask_pb2.FieldMask` :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() name = self.get_environment_name(project_id, region, environment_id) result = client.update_environment( request={"name": name, "environment": environment, "update_mask": update_mask}, retry=retry, timeout=timeout, metadata=metadata, ) return result
@GoogleBaseHook.fallback_to_default_project_id
[docs] def list_image_versions( self, project_id: str, region: str, page_size: int | None = None, page_token: str | None = None, include_past_releases: bool = False, retry: Retry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> ListImageVersionsPager: """ List ImageVersions for provided location. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param page_size: The maximum number of environments to return. :param page_token: The next_page_token value returned from a previous List request, if any. :param include_past_releases: Flag to include past releases :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_image_versions_client() result = client.list_image_versions( request={ "parent": self.get_parent(project_id, region), "page_size": page_size, "page_token": page_token, "include_past_releases": include_past_releases, }, retry=retry, timeout=timeout, metadata=metadata, ) return result
[docs]class CloudComposerAsyncHook(GoogleBaseHook): """Hook for Google Cloud Composer async APIs.""" def __init__(self, **kwargs): if kwargs.get("delegate_to") is not None: raise RuntimeError( "The `delegate_to` parameter has been deprecated before and finally removed in this version" " of Google Provider. You MUST convert it to `impersonate_chain`" ) super().__init__(**kwargs)
[docs] client_options = ClientOptions(api_endpoint="composer.googleapis.com:443")
[docs] def get_environment_client(self) -> EnvironmentsAsyncClient: """Retrieves client library object that allow access Environments service.""" return EnvironmentsAsyncClient( credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=self.client_options, )
[docs] def get_environment_name(self, project_id, region, environment_id): return f"projects/{project_id}/locations/{region}/environments/{environment_id}"
[docs] def get_parent(self, project_id, region): return f"projects/{project_id}/locations/{region}"
[docs] async def get_operation(self, operation_name): return await self.get_environment_client().transport.operations_client.get_operation( name=operation_name )
@GoogleBaseHook.fallback_to_default_project_id
[docs] async def create_environment( self, project_id: str, region: str, environment: Environment | dict, retry: AsyncRetry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> AsyncOperation: """ Create a new environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment: The environment to create. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() return await client.create_environment( request={"parent": self.get_parent(project_id, region), "environment": environment}, retry=retry, timeout=timeout, metadata=metadata, )
@GoogleBaseHook.fallback_to_default_project_id
[docs] async def delete_environment( self, project_id: str, region: str, environment_id: str, retry: AsyncRetry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> AsyncOperation: """ Delete an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() name = self.get_environment_name(project_id, region, environment_id) return await client.delete_environment( request={"name": name}, retry=retry, timeout=timeout, metadata=metadata )
@GoogleBaseHook.fallback_to_default_project_id
[docs] async def update_environment( self, project_id: str, region: str, environment_id: str, environment: Environment | dict, update_mask: dict | FieldMask, retry: AsyncRetry | _MethodDefault = DEFAULT, timeout: float | None = None, metadata: Sequence[tuple[str, str]] = (), ) -> AsyncOperation: r""" Update an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param environment: A patch environment. Fields specified by the ``updateMask`` will be copied from the patch environment into the environment under update. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param update_mask: Required. A comma-separated list of paths, relative to ``Environment``, of fields to update. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.protobuf.field_mask_pb2.FieldMask` :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """ client = self.get_environment_client() name = self.get_environment_name(project_id, region, environment_id) return await client.update_environment( request={"name": name, "environment": environment, "update_mask": update_mask}, retry=retry, timeout=timeout, metadata=metadata, )

Was this entry helpful?