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

Author: ca-nguyen

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

@@ -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]}")

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)

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)

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])