Source code for airflow.providers.amazon.aws.operators.step_function
# 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__importannotationsimportjsonfromdatetimeimporttimedeltafromtypingimportTYPE_CHECKING,Any,Sequencefromairflow.configurationimportconffromairflow.exceptionsimportAirflowExceptionfromairflow.providers.amazon.aws.hooks.step_functionimportStepFunctionHookfromairflow.providers.amazon.aws.links.step_functionimport(StateMachineDetailsLink,StateMachineExecutionsDetailsLink,)fromairflow.providers.amazon.aws.operators.base_awsimportAwsBaseOperatorfromairflow.providers.amazon.aws.triggers.step_functionimportStepFunctionsExecutionCompleteTriggerfromairflow.providers.amazon.aws.utilsimportvalidate_execute_complete_eventfromairflow.providers.amazon.aws.utils.mixinsimportaws_template_fieldsifTYPE_CHECKING:fromairflow.utils.contextimportContext
[docs]classStepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):""" An Operator that begins execution of an AWS Step Function State Machine. Additional arguments may be specified and are passed down to the underlying BaseOperator. .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:StepFunctionStartExecutionOperator` :param state_machine_arn: ARN of the Step Function State Machine :param name: The name of the execution. :param state_machine_input: JSON data input to pass to the State Machine :param aws_conn_id: The Airflow connection used for AWS credentials. If this is None or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param do_xcom_push: if True, execution_arn is pushed to XCom with key execution_arn. :param waiter_max_attempts: Maximum number of attempts to poll the execution. :param waiter_delay: Number of seconds between polling the state of the execution. :param deferrable: If True, the operator will wait asynchronously for the job to complete. This implies waiting for completion. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """
[docs]defexecute(self,context:Context):StateMachineDetailsLink.persist(context=context,operator=self,region_name=self.hook.conn_region_name,aws_partition=self.hook.conn_partition,state_machine_arn=self.state_machine_arn,)ifnot(execution_arn:=self.hook.start_execution(self.state_machine_arn,self.name,self.input)):raiseAirflowException(f"Failed to start State Machine execution for: {self.state_machine_arn}")StateMachineExecutionsDetailsLink.persist(context=context,operator=self,region_name=self.hook.conn_region_name,aws_partition=self.hook.conn_partition,execution_arn=execution_arn,)self.log.info("Started State Machine execution for %s: %s",self.state_machine_arn,execution_arn)ifself.deferrable:self.defer(trigger=StepFunctionsExecutionCompleteTrigger(execution_arn=execution_arn,waiter_delay=self.waiter_delay,waiter_max_attempts=self.waiter_max_attempts,aws_conn_id=self.aws_conn_id,region_name=self.region_name,botocore_config=self.botocore_config,verify=self.verify,),method_name="execute_complete",timeout=timedelta(seconds=self.waiter_max_attempts*self.waiter_delay),)returnexecution_arn
[docs]defexecute_complete(self,context:Context,event:dict[str,Any]|None=None)->None:event=validate_execute_complete_event(event)ifevent["status"]!="success":raiseAirflowException(f"Trigger error: event is {event}")self.log.info("State Machine execution completed successfully")returnevent["execution_arn"]
[docs]classStepFunctionGetExecutionOutputOperator(AwsBaseOperator[StepFunctionHook]):""" An Operator that returns the output of an AWS Step Function State Machine execution. Additional arguments may be specified and are passed down to the underlying BaseOperator. .. seealso:: For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:StepFunctionGetExecutionOutputOperator` :param execution_arn: ARN of the Step Function State Machine Execution :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """
[docs]defexecute(self,context:Context):StateMachineExecutionsDetailsLink.persist(context=context,operator=self,region_name=self.hook.conn_region_name,aws_partition=self.hook.conn_partition,execution_arn=self.execution_arn,)execution_status=self.hook.describe_execution(self.execution_arn)response=Noneif"output"inexecution_status:response=json.loads(execution_status["output"])elif"error"inexecution_status:response=json.loads(execution_status["error"])self.log.info("Got State Machine Execution output for %s",self.execution_arn)returnresponse