Source code for airflow.providers.google.cloud.sensors.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__future__importannotationsfromcollections.abcimportSequencefromtypingimportTYPE_CHECKINGfromgoogle.api_core.gapic_v1.methodimportDEFAULT,_MethodDefaultfromgoogle.cloud.workflows.executions_v1betaimportExecutionfromairflow.exceptionsimportAirflowExceptionfromairflow.providers.google.cloud.hooks.workflowsimportWorkflowsHookfromairflow.providers.google.common.hooks.base_googleimportPROVIDE_PROJECT_IDfromairflow.sensors.baseimportBaseSensorOperatorifTYPE_CHECKING:fromgoogle.api_core.retryimportRetryfromairflow.utils.contextimportContext
[docs]classWorkflowExecutionSensor(BaseSensorOperator):""" Checks state of an execution for the given ``workflow_id`` and ``execution_id``. :param workflow_id: Required. The ID of the workflow. :param execution_id: Required. The ID of the execution. :param project_id: Required. The ID of the Google Cloud project the cluster belongs to. :param location: Required. The Cloud Dataproc region in which to handle the request. :param success_states: Execution states to be considered as successful, by default it's only ``SUCCEEDED`` state :param failure_states: Execution states to be considered as failures, by default they are ``FAILED`` and ``CANCELLED`` states. :param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be retried. :param request_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. :param metadata: Additional metadata that is provided to the method. """
[docs]defpoke(self,context:Context):hook=WorkflowsHook(gcp_conn_id=self.gcp_conn_id,impersonation_chain=self.impersonation_chain)self.log.info("Checking state of execution %s for workflow %s",self.execution_id,self.workflow_id)execution:Execution=hook.get_execution(workflow_id=self.workflow_id,execution_id=self.execution_id,location=self.location,project_id=self.project_id,retry=self.retry,timeout=self.request_timeout,metadata=self.metadata,)state=execution.stateifstateinself.failure_states:message=(f"Execution {self.execution_id} for workflow {self.execution_id} "f"failed and is in `{state}` state")raiseAirflowException(message)ifstateinself.success_states:self.log.info("Execution %s for workflow %s completed with state: %s",self.execution_id,self.workflow_id,state,)returnTrueself.log.info("Execution %s for workflow %s does not completed yet, current state: %s",self.execution_id,self.workflow_id,state,)returnFalse