Add RedriveExecution support to StepFunctionStartExecutionOperator

airflow/providers/amazon/aws/hooks/step_function.py

@@ -18,6 +18,7 @@
 
 import json
 
+from airflow.exceptions import AirflowFailException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
 
@@ -43,6 +44,7 @@ def start_execution(
         state_machine_arn: str,
         name: str | None = None,
         state_machine_input: dict | str | None = None,
+        is_redrive_execution: bool = False,
     ) -> str:
         """
         Start Execution of the State Machine.
@@ -51,10 +53,26 @@ def start_execution(
             - :external+boto3:py:meth:`SFN.Client.start_execution`
 
         :param state_machine_arn: AWS Step Function State Machine ARN.
+        :param is_redrive_execution: Restarts unsuccessful executions of Standard workflows that did not
+            complete successfully in the last 14 days.
         :param name: The name of the execution.
         :param state_machine_input: JSON data input to pass to the State Machine.
         :return: Execution ARN.
         """
+        if is_redrive_execution:
+            if not name:
+                raise AirflowFailException(
+                    "Execution name is required to start RedriveExecution for %s.", state_machine_arn
+                )
+            elements = state_machine_arn.split(":stateMachine:")
+            execution_arn = f"{elements[0]}:execution:{elements[1]}:{name}"
+            self.conn.redrive_execution(executionArn=execution_arn)
+            self.log.info(
+                "Successfully started RedriveExecution for Step Function State Machine: %s.",
+                state_machine_arn,
+            )
+            return execution_arn
+
         execution_args = {"stateMachineArn": state_machine_arn}
         if name is not None:
             execution_args["name"] = name

airflow/providers/amazon/aws/operators/step_function.py

@@ -48,6 +48,8 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
 
     :param state_machine_arn: ARN of the Step Function State Machine
     :param name: The name of the execution.
+    :param is_redrive_execution: Restarts unsuccessful executions of Standard workflows that did not
+            complete successfully in the last 14 days.
     :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
@@ -73,7 +75,9 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
     """
 
     aws_hook_class = StepFunctionHook
-    template_fields: Sequence[str] = aws_template_fields("state_machine_arn", "name", "input")
+    template_fields: Sequence[str] = aws_template_fields(
+        "state_machine_arn", "name", "input", "is_redrive_execution"
+    )
     ui_color = "#f9c915"
     operator_extra_links = (StateMachineDetailsLink(), StateMachineExecutionsDetailsLink())
 
@@ -82,6 +86,7 @@ def __init__(
         *,
         state_machine_arn: str,
         name: str | None = None,
+        is_redrive_execution: bool = False,
         state_machine_input: dict | str | None = None,
         waiter_max_attempts: int = 30,
         waiter_delay: int = 60,
@@ -91,6 +96,7 @@ def __init__(
         super().__init__(**kwargs)
         self.state_machine_arn = state_machine_arn
         self.name = name
+        self.is_redrive_execution = is_redrive_execution
         self.input = state_machine_input
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
@@ -105,7 +111,11 @@ def execute(self, context: Context):
             state_machine_arn=self.state_machine_arn,
         )
 
-        if not (execution_arn := self.hook.start_execution(self.state_machine_arn, self.name, self.input)):
+        if not (
+            execution_arn := self.hook.start_execution(
+                self.state_machine_arn, self.name, self.input, self.is_redrive_execution
+            )
+        ):
             raise AirflowException(f"Failed to start State Machine execution for: {self.state_machine_arn}")
 
         StateMachineExecutionsDetailsLink.persist(

tests/providers/amazon/aws/hooks/test_step_function.py

@@ -17,8 +17,13 @@
 # under the License.
 from __future__ import annotations
 
+from datetime import datetime
+from unittest import mock
+
+import pytest
 from moto import mock_aws
 
+from airflow.exceptions import AirflowFailException
 from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
 
 
@@ -42,6 +47,31 @@ def test_start_execution(self):
 
         assert execution_arn is not None
 
+    @mock.patch.object(StepFunctionHook, "conn")
+    def test_redrive_execution(self, mock_conn):
+        mock_conn.redrive_execution.return_value = {"redriveDate": datetime(2024, 1, 1)}
+        StepFunctionHook().start_execution(
+            state_machine_arn="arn:aws:states:us-east-1:123456789012:stateMachine:test-state-machine",
+            name="random-123",
+            is_redrive_execution=True,
+        )
+
+        mock_conn.redrive_execution.assert_called_once_with(
+            executionArn="arn:aws:states:us-east-1:123456789012:execution:test-state-machine:random-123"
+        )
+
+    @mock.patch.object(StepFunctionHook, "conn")
+    def test_redrive_execution_without_name_should_fail(self, mock_conn):
+        mock_conn.redrive_execution.return_value = {"redriveDate": datetime(2024, 1, 1)}
+
+        with pytest.raises(
+            AirflowFailException, match="Execution name is required to start RedriveExecution"
+        ):
+            StepFunctionHook().start_execution(
+                state_machine_arn="arn:aws:states:us-east-1:123456789012:stateMachine:test-state-machine",
+                is_redrive_execution=True,
+            )
+
     def test_describe_execution(self):
         hook = StepFunctionHook(aws_conn_id="aws_default", region_name="us-east-1")
         state_machine = hook.get_conn().create_state_machine(

tests/providers/amazon/aws/operators/test_step_function.py

@@ -159,7 +159,7 @@ def test_execute(self, mocked_hook, mocked_context):
             aws_conn_id=None,
         )
         assert op.execute(mocked_context) == hook_response
-        mocked_hook.start_execution.assert_called_once_with(STATE_MACHINE_ARN, NAME, INPUT)
+        mocked_hook.start_execution.assert_called_once_with(STATE_MACHINE_ARN, NAME, INPUT, False)
         self.mocked_details_link.assert_called_once_with(
             aws_partition=mock.ANY,
             context=mock.ANY,
@@ -189,7 +189,7 @@ def test_step_function_start_execution_deferrable(self, mocked_hook):
         )
         with pytest.raises(TaskDeferred):
             operator.execute(None)
-        mocked_hook.start_execution.assert_called_once_with(STATE_MACHINE_ARN, NAME, INPUT)
+        mocked_hook.start_execution.assert_called_once_with(STATE_MACHINE_ARN, NAME, INPUT, False)
 
     @mock.patch.object(StepFunctionStartExecutionOperator, "hook")
     @pytest.mark.parametrize("execution_arn", [pytest.param(None, id="none"), pytest.param("", id="empty")])
@@ -200,3 +200,35 @@ def test_step_function_no_execution_arn_returns(self, mocked_hook, execution_arn
         )
         with pytest.raises(AirflowException, match="Failed to start State Machine execution"):
             op.execute({})
+
+    @mock.patch.object(StepFunctionStartExecutionOperator, "hook")
+    def test_start_redrive_execution(self, mocked_hook, mocked_context):
+        hook_response = (
+            "arn:aws:states:us-east-1:123456789012:execution:"
+            "pseudo-state-machine:020f5b16-b1a1-4149-946f-92dd32d97934"
+        )
+        mocked_hook.start_execution.return_value = hook_response
+        op = StepFunctionStartExecutionOperator(
+            task_id=self.TASK_ID,
+            state_machine_arn=STATE_MACHINE_ARN,
+            name=NAME,
+            is_redrive_execution=True,
+            state_machine_input=None,
+            aws_conn_id=None,
+        )
+        assert op.execute(mocked_context) == hook_response
+        mocked_hook.start_execution.assert_called_once_with(STATE_MACHINE_ARN, NAME, None, True)
+        self.mocked_details_link.assert_called_once_with(
+            aws_partition=mock.ANY,
+            context=mock.ANY,
+            operator=mock.ANY,
+            region_name=mock.ANY,
+            state_machine_arn=STATE_MACHINE_ARN,
+        )
+        self.mocked_executions_details_link.assert_called_once_with(
+            aws_partition=mock.ANY,
+            context=mock.ANY,
+            operator=mock.ANY,
+            region_name=mock.ANY,
+            execution_arn=EXECUTION_ARN,
+        )