feat: Add support for nested Aws StepFunctions service integration

doc/services.rst

@@ -20,6 +20,8 @@ This module provides classes to build steps that integrate with Amazon DynamoDB,
 
 - `Amazon SQS <#amazon-sqs>`__
 
+- `AWS Step Functions <#aws-step-functions>`__
+
 
 Amazon DynamoDB
 ----------------
@@ -82,3 +84,8 @@ Amazon SNS
 Amazon SQS
 -----------
 .. autoclass:: stepfunctions.steps.service.SqsSendMessageStep
+
+AWS Step Functions
+------------------
+.. autoclass:: stepfunctions.steps.service.StepFunctionsStartExecutionStep
+

src/stepfunctions/steps/__init__.py

@@ -34,3 +34,4 @@
 from stepfunctions.steps.service import EventBridgePutEventsStep
 from stepfunctions.steps.service import GlueDataBrewStartJobRunStep
 from stepfunctions.steps.service import SnsPublishStep, SqsSendMessageStep
+from stepfunctions.steps.service import StepFunctionsStartExecutionStep

src/stepfunctions/steps/integration_resources.py

@@ -13,9 +13,13 @@
 
 from __future__ import absolute_import
 
+import logging
+
 from enum import Enum
 from stepfunctions.steps.utils import get_aws_partition
 
+logger = logging.getLogger('stepfunctions')
+
 
 class IntegrationPattern(Enum):
     """
@@ -25,6 +29,17 @@ class IntegrationPattern(Enum):
     WaitForTaskToken = "waitForTaskToken"
     WaitForCompletion = "sync"
     RequestResponse = ""
+    WaitForCompletionWithJsonResponse = "sync:2"
+
+
+class ServiceIntegrationType(Enum):
+    """
+    Service Integration Types for service integration resources (see `Service Integration Patterns <https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html`_ for more details.)
+    """
+
+    REQUEST_RESPONSE = "RequestResponse",
+    RUN_A_JOB = "RunAJob",
+    WAIT_FOR_CALLBACK = "WaitForCallback"
 
 
 def get_service_integration_arn(service, api, integration_pattern=IntegrationPattern.RequestResponse):
@@ -44,3 +59,20 @@ def get_service_integration_arn(service, api, integration_pattern=IntegrationPat
     return arn
 
 
+def get_integration_pattern_from_service_integration_type(service_integration_type):
+    """
+    Returns the integration pattern to use for the service integration type.
+    IntegrationPattern.RequestResponse is returned by default if the service type is not recognized
+    Args:
+        service_integration_type(ServiceIntegrationType): Service integration type to use to get the integration pattern
+    """
+
+    if service_integration_type == ServiceIntegrationType.RUN_A_JOB:
+        return IntegrationPattern.WaitForCompletion
+    elif service_integration_type == ServiceIntegrationType.WAIT_FOR_CALLBACK:
+        return IntegrationPattern.WaitForTaskToken
+    else:
+        if not isinstance(service_integration_type, ServiceIntegrationType):
+            logger.warning(f"Invalid Service integration type ({service_integration_type}) - returning IntegrationPattern.RequestResponse by default")
+        return IntegrationPattern.RequestResponse
+

src/stepfunctions/steps/service.py

@@ -15,7 +15,8 @@
 from enum import Enum
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
-from stepfunctions.steps.integration_resources import IntegrationPattern, get_service_integration_arn
+from stepfunctions.steps.integration_resources import IntegrationPattern, ServiceIntegrationType,\
+    get_integration_pattern_from_service_integration_type, get_service_integration_arn
 
 DYNAMODB_SERVICE_NAME = "dynamodb"
 EKS_SERVICES_NAME = "eks"
@@ -24,6 +25,7 @@
 GLUE_DATABREW_SERVICE_NAME = "databrew"
 SNS_SERVICE_NAME = "sns"
 SQS_SERVICE_NAME = "sqs"
+STEP_FUNCTIONS_SERVICE_NAME = "states"
 
 
 class DynamoDBApi(Enum):
@@ -70,6 +72,10 @@ class SqsApi(Enum):
     SendMessage = "sendMessage"
 
 
+class StepFunctions(Enum):
+    StartExecution = "startExecution"
+
+
 class DynamoDBGetItemStep(Task):
     """
     Creates a Task state to get an item from DynamoDB. See `Call DynamoDB APIs with Step Functions <https://docs.aws.amazon.com/step-functions/latest/dg/connect-ddb.html>`_ for more details.
@@ -887,3 +893,60 @@ def __init__(self, state_id, **kwargs):
                                                                    ElasticMapReduceApi.ModifyInstanceGroupByName)
 
         super(EmrModifyInstanceGroupByNameStep, self).__init__(state_id, **kwargs)
+
+
+class StepFunctionsStartExecutionStep(Task):
+
+    """
+    Creates a Task state that starts an execution of another state machine. See `Manage AWS Step Functions Executions as an Integrated Service <https://docs.aws.amazon.com/step-functions/latest/dg/connect-stepfunctions.html`_ for more details.
+
+    Property flags: There are three flags (wait_for_callback, wait_for_completion and async_call) that can be set in order to select which Step Functions resource to use.
+        One of three must be enabled to create the step successfully.
+    """
+
+    def __init__(self, state_id, service_integration_type, **kwargs):
+        """
+        Args:
+            state_id (str): State name whose length **must be** less than or equal to 128 unicode characters. State names **must be** unique within the scope of the whole state machine.
+            integration_pattern (stepfunctions.steps.integration_resources.ServiceIntegrationType): Service integration type to use to build resource.
+                Supported service integration types: REQUEST_RESPONSE, RUN_A_JOB and WAIT_FOR_CALLBACK
+            timeout_seconds (int, optional): Positive integer specifying timeout for the state in seconds. If the state runs longer than the specified timeout, then the interpreter fails the state with a `States.Timeout` Error Name. (default: 60)
+            timeout_seconds_path (str, optional): Path specifying the state's timeout value in seconds from the state input. When resolved, the path must select a field whose value is a positive integer.
+            heartbeat_seconds (int, optional): Positive integer specifying heartbeat timeout for the state in seconds. This value should be lower than the one specified for `timeout_seconds`. If more time than the specified heartbeat elapses between heartbeats from the task, then the interpreter fails the state with a `States.Timeout` Error Name.
+            heartbeat_seconds_path (str, optional): Path specifying the state's heartbeat value in seconds from the state input. When resolved, the path must select a field whose value is a positive integer.
+            comment (str, optional): Human-readable comment or description. (default: None)
+            input_path (str, optional): Path applied to the state’s raw input to select some or all of it; that selection is used by the state. (default: '$')
+            parameters (dict, optional): The value of this field becomes the effective input for the state.
+            result_path (str, optional): Path specifying the raw input’s combination with or replacement by the state’s result. (default: '$')
+            output_path (str, optional): Path applied to the state’s output after the application of `result_path`, producing the effective output which serves as the raw input for the next state. (default: '$')
+        """
+        supported_integ_types = [ServiceIntegrationType.REQUEST_RESPONSE, ServiceIntegrationType.RUN_A_JOB,
+                                 ServiceIntegrationType.WAIT_FOR_CALLBACK]
+
+        if not isinstance(service_integration_type, ServiceIntegrationType):
+            raise ValueError(f"Invalid type used for service_integration_type arg ({service_integration_type}, "
+                             f"{type(service_integration_type)}). Accepted type: {ServiceIntegrationType}")
+        elif service_integration_type not in supported_integ_types:
+            raise ValueError(f"Service Integration Type ({service_integration_type.name}) is not supported for this step - "
+                             f"Please use one of the following: "
+                             f"{[integ_type.name for integ_type in supported_integ_types]}")
+
+        if service_integration_type == ServiceIntegrationType.RUN_A_JOB:
+            # For RunAJob service integration type, use the resource that returns json response (`.sync:2`)
+            """
+            Example resource arn:aws:states:::states:startExecution.sync:2
+            """
+            kwargs[Field.Resource.value] = get_service_integration_arn(STEP_FUNCTIONS_SERVICE_NAME,
+                                                                       StepFunctions.StartExecution,
+                                                                       IntegrationPattern.WaitForCompletionWithJsonResponse)
+        else:
+            """
+            Example resource arn:
+                - arn:aws:states:::states:startExecution.waitForTaskToken
+                - arn:aws:states:::states:startExecution
+            """
+            kwargs[Field.Resource.value] = get_service_integration_arn(STEP_FUNCTIONS_SERVICE_NAME,
+                                                                       StepFunctions.StartExecution,
+                                                                       get_integration_pattern_from_service_integration_type(service_integration_type))
+
+        super(StepFunctionsStartExecutionStep, self).__init__(state_id, **kwargs)

tests/unit/test_service_steps.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import boto3
+import pytest
 
 from unittest.mock import patch
 from stepfunctions.steps.service import DynamoDBGetItemStep, DynamoDBPutItemStep, DynamoDBUpdateItemStep, DynamoDBDeleteItemStep
@@ -30,6 +31,8 @@
 from stepfunctions.steps.service import EventBridgePutEventsStep
 from stepfunctions.steps.service import SnsPublishStep, SqsSendMessageStep
 from stepfunctions.steps.service import GlueDataBrewStartJobRunStep
+from stepfunctions.steps.service import StepFunctionsStartExecutionStep
+from stepfunctions.steps.integration_resources import IntegrationPattern, ServiceIntegrationType
 
 
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
@@ -1158,3 +1161,92 @@ def test_eks_call_step_creation():
         },
         'End': True
     }
+
+
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_step_functions_start_execution_step_creation():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution", service_integration_type=ServiceIntegrationType.REQUEST_RESPONSE, parameters={
+            "Input": {
+                "Comment": "Hello world!"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        })
+
+    assert step.to_dict() == {
+        "Type": "Task",
+        "Resource": "arn:aws:states:::states:startExecution",
+        "Parameters": {
+            "Input": {
+                "Comment": "Hello world!"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        },
+        "End": True
+    }
+
+
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_step_functions_start_execution_step_creation_sync():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution - Sync", service_integration_type=ServiceIntegrationType.RUN_A_JOB, parameters={
+            "Input": {
+                "Comment": "Hello world!"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        })
+
+    assert step.to_dict() == {
+        "Type": "Task",
+        "Resource": "arn:aws:states:::states:startExecution.sync:2",
+        "Parameters": {
+            "Input": {
+                "Comment": "Hello world!"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        },
+        "End": True
+    }
+
+
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_step_functions_start_execution_step_creation_wait_for_callback():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution - Wait for Callback", service_integration_type=ServiceIntegrationType.WAIT_FOR_CALLBACK,
+        parameters={
+            "Input": {
+                "Comment": "Hello world!",
+                "token.$": "$$.Task.Token"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        })
+
+    assert step.to_dict() == {
+        "Type": "Task",
+        "Resource": "arn:aws:states:::states:startExecution.waitForTaskToken",
+        "Parameters": {
+            "Input": {
+                "Comment": "Hello world!",
+                "token.$": "$$.Task.Token"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        },
+        "End": True
+    }
+
+
+@pytest.mark.parametrize("service_integration_type", [
+    None,
+    "ServiceIntegrationTypeStr",
+    IntegrationPattern.RequestResponse
+])
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_step_functions_start_execution_step_creation_invalid_service_integration_type_raises_exception(service_integration_type):
+    with pytest.raises(ValueError):
+        StepFunctionsStartExecutionStep("SFN Start Execution - invalid ServiceType", service_integration_type=service_integration_type)

tests/unit/test_steps_utils.py

@@ -13,11 +13,16 @@
 
 # Test if boto3 session can fetch correct aws partition info from test environment
 
-from stepfunctions.steps.utils import get_aws_partition, merge_dicts
-from stepfunctions.steps.integration_resources import IntegrationPattern, get_service_integration_arn
 import boto3
-from unittest.mock import patch
+import logging
+import pytest
+
 from enum import Enum
+from unittest.mock import patch
+
+from stepfunctions.steps.utils import get_aws_partition, merge_dicts
+from stepfunctions.steps.integration_resources import IntegrationPattern, ServiceIntegrationType,\
+    get_integration_pattern_from_service_integration_type, get_service_integration_arn
 
 
 testService = "sagemaker"
@@ -86,3 +91,27 @@ def test_merge_dicts():
         'b': 2,
         'c': 3
     }
+
+
+@pytest.mark.parametrize("service_integration_type, expected_integration_pattern", [
+    (ServiceIntegrationType.REQUEST_RESPONSE, IntegrationPattern.RequestResponse),
+    (ServiceIntegrationType.RUN_A_JOB, IntegrationPattern.WaitForCompletion),
+    (ServiceIntegrationType.WAIT_FOR_CALLBACK, IntegrationPattern.WaitForTaskToken)
+])
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_get_integration_pattern_from_service_integration_type(service_integration_type, expected_integration_pattern):
+    integration_pattern = get_integration_pattern_from_service_integration_type(service_integration_type)
+    assert integration_pattern == expected_integration_pattern
+
+
+@pytest.mark.parametrize("service_integration_type", [
+    None,
+    "ServiceIntegrationTypeStr",
+    IntegrationPattern.RequestResponse
+])
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_get_integration_pattern_from_service_integration_type_with_invalid_type(service_integration_type, caplog):
+    with caplog.at_level(logging.WARNING):
+        integration_pattern = get_integration_pattern_from_service_integration_type(service_integration_type)
+        assert 'WARNING' in caplog.text
+    assert integration_pattern == IntegrationPattern.RequestResponse