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

# 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 typing import Dict, Optional, Sequence, Tuple, Union

from google.api_core.operation import Operation
from google.api_core.retry import Retry

# pylint: disable=no-name-in-module
from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient
from google.cloud.workflows.executions_v1beta.services.executions.pagers import ListExecutionsPager
from google.cloud.workflows_v1beta import Workflow, WorkflowsClient
from google.cloud.workflows_v1beta.services.workflows.pagers import ListWorkflowsPager
from google.protobuf.field_mask_pb2 import FieldMask

from airflow.providers.google.common.hooks.base_google import GoogleBaseHook

# pylint: enable=no-name-in-module


[docs]class WorkflowsHook(GoogleBaseHook): """ Hook for Google GCP APIs. All the methods in the hook where project_id is used must be called with keyword arguments rather than positional. """
[docs] def get_workflows_client(self) -> WorkflowsClient: """Returns WorkflowsClient.""" return WorkflowsClient(credentials=self._get_credentials(), client_info=self.client_info)
[docs] def get_executions_client(self) -> ExecutionsClient: """Returns ExecutionsClient.""" return ExecutionsClient(credentials=self._get_credentials(), client_info=self.client_info)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def create_workflow( self, workflow: Dict, workflow_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Operation: """ Creates a new workflow. If a workflow with the specified name already exists in the specified project and location, the long running operation will return [ALREADY_EXISTS][google.rpc.Code.ALREADY_EXISTS] error. :param workflow: Required. Workflow to be created. :type workflow: Dict :param workflow_id: Required. The ID of the workflow to be created. :type workflow_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_workflows_client() parent = f"projects/{project_id}/locations/{location}" return client.create_workflow( request={"parent": parent, "workflow": workflow, "workflow_id": workflow_id}, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleBaseHook.fallback_to_default_project_id
[docs] def get_workflow( self, workflow_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Workflow: """ Gets details of a single Workflow. :param workflow_id: Required. The ID of the workflow to be created. :type workflow_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_workflows_client() name = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}" return client.get_workflow(request={"name": name}, retry=retry, timeout=timeout, metadata=metadata)
[docs] def update_workflow( self, workflow: Union[Dict, Workflow], update_mask: Optional[FieldMask] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Operation: """ Updates an existing workflow. Running this method has no impact on already running executions of the workflow. A new revision of the workflow may be created as a result of a successful update operation. In that case, such revision will be used in new workflow executions. :param workflow: Required. Workflow to be created. :type workflow: Dict :param update_mask: List of fields to be updated. If not present, the entire workflow will be updated. :type update_mask: FieldMask :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_workflows_client() return client.update_workflow( request={"workflow": workflow, "update_mask": update_mask}, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleBaseHook.fallback_to_default_project_id
[docs] def delete_workflow( self, workflow_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Operation: """ Deletes a workflow with the specified name. This method also cancels and deletes all running executions of the workflow. :param workflow_id: Required. The ID of the workflow to be created. :type workflow_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_workflows_client() name = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}" return client.delete_workflow(request={"name": name}, retry=retry, timeout=timeout, metadata=metadata)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def list_workflows( self, location: str, project_id: str, filter_: Optional[str] = None, order_by: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> ListWorkflowsPager: """ Lists Workflows in a given project and location. The default order is not specified. :param filter_: Filter to restrict results to specific workflows. :type filter_: str :param order_by: Comma-separated list of fields that that specify the order of the results. Default sorting order for a field is ascending. To specify descending order for a field, append a "desc" suffix. If not specified, the results will be returned in an unspecified order. :type order_by: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_workflows_client() parent = f"projects/{project_id}/locations/{location}" return client.list_workflows( request={"parent": parent, "filter": filter_, "order_by": order_by}, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleBaseHook.fallback_to_default_project_id
[docs] def create_execution( self, workflow_id: str, location: str, project_id: str, execution: Dict, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Execution: """ Creates a new execution using the latest revision of the given workflow. :param execution: Required. Input parameters of the execution represented as a dictionary. :type execution: Dict :param workflow_id: Required. The ID of the workflow. :type workflow_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_executions_client() parent = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}" return client.create_execution( request={"parent": parent, "execution": execution}, retry=retry, timeout=timeout, metadata=metadata,
) @GoogleBaseHook.fallback_to_default_project_id
[docs] def get_execution( self, workflow_id: str, execution_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Execution: """ Returns an execution for the given ``workflow_id`` and ``execution_id``. :param workflow_id: Required. The ID of the workflow. :type workflow_id: str :param execution_id: Required. The ID of the execution. :type execution_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_executions_client() name = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}/executions/{execution_id}" return client.get_execution(request={"name": name}, retry=retry, timeout=timeout, metadata=metadata)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def cancel_execution( self, workflow_id: str, execution_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Execution: """ Cancels an execution using the given ``workflow_id`` and ``execution_id``. :param workflow_id: Required. The ID of the workflow. :type workflow_id: str :param execution_id: Required. The ID of the execution. :type execution_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_executions_client() name = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}/executions/{execution_id}" return client.cancel_execution( request={"name": name}, retry=retry, timeout=timeout, metadata=metadata
) @GoogleBaseHook.fallback_to_default_project_id
[docs] def list_executions( self, workflow_id: str, location: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> ListExecutionsPager: """ Returns a list of executions which belong to the workflow with the given name. The method returns executions of all workflow revisions. Returned executions are ordered by their start time (newest first). :param workflow_id: Required. The ID of the workflow to be created. :type workflow_id: str :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :type project_id: str :param location: Required. The GCP region in which to handle the request. :type location: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :type retry: google.api_core.retry.Retry :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if ``retry`` is specified, the timeout applies to each individual attempt. :type timeout: float :param metadata: Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]] """ metadata = metadata or () client = self.get_executions_client() parent = f"projects/{project_id}/locations/{location}/workflows/{workflow_id}" return client.list_executions( request={"parent": parent}, retry=retry, timeout=timeout, metadata=metadata
)

Was this entry helpful?