feat: Add support for nested Aws StepFunctions service integration (#166)

ca-nguyen · web-flow · commit f9e63b312f54 · 2021-10-21T15:25:21.000-07:00
diff --git a/doc/services.rst b/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
+
diff --git a/src/stepfunctions/steps/__init__.py b/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
diff --git a/src/stepfunctions/steps/integration_resources.py b/src/stepfunctions/steps/integration_resources.py
@@ -24,23 +24,35 @@ class IntegrationPattern(Enum):
 
     WaitForTaskToken = "waitForTaskToken"
     WaitForCompletion = "sync"
-    RequestResponse = ""
+    CallAndContinue = ""
 
 
-def get_service_integration_arn(service, api, integration_pattern=IntegrationPattern.RequestResponse):
+def get_service_integration_arn(service, api, integration_pattern=IntegrationPattern.CallAndContinue, version=None):
 
     """
     ARN builder for task integration
     Args:
         service (str): The service name for the service integration
         api (str): The api of the service integration
-        integration_pattern (IntegrationPattern, optional): The integration pattern for the task. (Default: IntegrationPattern.RequestResponse)
+        integration_pattern (IntegrationPattern, optional): The integration pattern for the task. (Default: IntegrationPattern.CallAndContinue)
+        version (int, optional): The version of the resource to use. (Default: None)
     """
     arn = ""
-    if integration_pattern == IntegrationPattern.RequestResponse:
+    if integration_pattern == IntegrationPattern.CallAndContinue:
         arn = f"arn:{get_aws_partition()}:states:::{service}:{api.value}"
     else:
         arn = f"arn:{get_aws_partition()}:states:::{service}:{api.value}.{integration_pattern.value}"
+
+    if version:
+        arn = f"{arn}:{str(version)}"
+
     return arn
 
 
+def is_integration_pattern_valid(integration_pattern, supported_integration_patterns):
+    if not isinstance(integration_pattern, IntegrationPattern):
+        raise TypeError(f"Integration pattern must be of type {IntegrationPattern}")
+    elif integration_pattern not in supported_integration_patterns:
+        raise ValueError(f"Integration Pattern ({integration_pattern.name}) is not supported for this step - "
+                         f"Please use one of the following: "
+                         f"{[integ_type.name for integ_type in supported_integration_patterns]}")
diff --git a/src/stepfunctions/steps/service.py b/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, get_service_integration_arn,\
+    is_integration_pattern_valid
 
 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,54 @@ 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 a 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.
+    """
+
+    def __init__(self, state_id, integration_pattern=IntegrationPattern.WaitForCompletion, **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.IntegrationPattern, optional): Service integration pattern used to call the integrated service. (default: WaitForCompletion)
+                Supported integration patterns:
+                    WaitForCompletion: Wait for the state machine execution to complete before going to the next state. (See `Run A Job <https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-sync`_ for more details.)
+                    WaitForTaskToken: Wait for the state machine execution to return a task token before progressing to the next state (See `Wait for a Callback with the Task Token <https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-wait-token`_ for more details.)
+                    CallAndContinue: Call StartExecution and progress to the next state (See `Request Response <https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-default`_ for more details.)
+            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. (default: None)
+            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_patterns = [IntegrationPattern.WaitForCompletion, IntegrationPattern.WaitForTaskToken,
+                                    IntegrationPattern.CallAndContinue]
+
+        is_integration_pattern_valid(integration_pattern, supported_integ_patterns)
+
+        if integration_pattern == IntegrationPattern.WaitForCompletion:
+            """
+            Example resource arn:aws:states:::states:startExecution.sync:2
+            """
+            kwargs[Field.Resource.value] = get_service_integration_arn(STEP_FUNCTIONS_SERVICE_NAME,
+                                                                       StepFunctions.StartExecution,
+                                                                       integration_pattern,
+                                                                       2)
+        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,
+                                                                       integration_pattern)
+
+        super(StepFunctionsStartExecutionStep, self).__init__(state_id, **kwargs)
diff --git a/tests/unit/test_service_steps.py b/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
 
 
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
@@ -1158,3 +1161,97 @@ 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_default():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution", parameters={
+            "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": {
+            "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_call_and_continue():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution", integration_pattern=IntegrationPattern.CallAndContinue, parameters={
+            "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": {
+            "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_completion():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution - Sync", integration_pattern=IntegrationPattern.WaitForCompletion, parameters={
+            "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": {
+            "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_task_token():
+    step = StepFunctionsStartExecutionStep(
+        "SFN Start Execution - Wait for Callback", integration_pattern=IntegrationPattern.WaitForTaskToken,
+        parameters={
+            "Input": {
+                "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": {
+                "token.$": "$$.Task.Token"
+            },
+            "StateMachineArn": "arn:aws:states:us-east-1:123456789012:stateMachine:HelloWorld",
+            "Name": "ExecutionName"
+        },
+        "End": True
+    }
+
+
+@pytest.mark.parametrize("integration_pattern", [
+    None,
+    "ServiceIntegrationTypeStr",
+    0
+])
+@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
+def test_step_functions_start_execution_step_creation_invalid_integration_pattern_raises_type_error(integration_pattern):
+    with pytest.raises(TypeError):
+        StepFunctionsStartExecutionStep("SFN Start Execution - invalid ServiceType", integration_pattern=integration_pattern)
diff --git a/tests/unit/test_steps_utils.py b/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,\
+    get_service_integration_arn, is_integration_pattern_valid
 
 
 testService = "sagemaker"
@@ -86,3 +91,18 @@ def test_merge_dicts():
         'b': 2,
         'c': 3
     }
+
+
+@pytest.mark.parametrize("service_integration_type", [
+    None,
+    "IntegrationPatternStr",
+    0
+])
+def test_is_integration_pattern_valid_with_invalid_type_raises_type_error(service_integration_type):
+    with pytest.raises(TypeError):
+        is_integration_pattern_valid(service_integration_type, [IntegrationPattern.WaitForTaskToken])
+
+
+def test_is_integration_pattern_valid_with_non_supported_type_raises_value_error():
+    with pytest.raises(ValueError):
+        is_integration_pattern_valid(IntegrationPattern.WaitForTaskToken, [IntegrationPattern.WaitForCompletion])
