def get_hook(self) -> StepFunctionHook:
"""Create and return a StepFunctionHook"""
if self.hook:
return self.hook
self.hook = StepFunctionHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
return self.hook
def execute(self, context: 'Context'):
hook = StepFunctionHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
execution_status = hook.describe_execution(self.execution_arn)
execution_output = json.loads(execution_status['output']) if 'output' in execution_status else None
self.log.info('Got State Machine Execution output for %s', self.execution_arn)
return execution_output
def execute(self, context: 'Context'):
hook = StepFunctionHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
execution_arn = hook.start_execution(self.state_machine_arn, self.name, self.input)
if execution_arn is None:
raise AirflowException(f'Failed to start State Machine execution for: {self.state_machine_arn}')
self.log.info('Started State Machine execution for %s: %s', self.state_machine_arn, execution_arn)
return execution_arn
def test_start_execution(self):
hook = StepFunctionHook(aws_conn_id='aws_default',
region_name='us-east-1')
state_machine = hook.get_conn().create_state_machine(
name='pseudo-state-machine',
definition='{}',
roleArn='arn:aws:iam::000000000000:role/Role')
state_machine_arn = state_machine.get('stateMachineArn', None)
execution_arn = hook.start_execution(
state_machine_arn=state_machine_arn,
name=None,
state_machine_input={})
assert execution_arn is not None
def test_get_conn_returns_a_boto3_connection(self):
hook = StepFunctionHook(aws_conn_id='aws_default')
self.assertEqual('stepfunctions',
hook.get_conn().meta.service_model.service_name)