fix: update arn builder method and its usage

Author: Bowen Yuan

src/stepfunctions/steps/compute.py

@@ -12,10 +12,32 @@
 # permissions and limitations under the License.
 from __future__ import absolute_import
 
+from enum import Enum
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
-from stepfunctions.steps.utils import resource_integration_arn_builder
-from stepfunctions.steps.integration_resources import IntegrationPattern, IntegrationServices, LambdaApi, GlueApi, BatchApi, EcsApi
+from stepfunctions.steps.utils import get_service_integration_arn
+from stepfunctions.steps.integration_resources import IntegrationPattern
+
+Lambda = "lambda"
+Glue = "glue"
+Ecs = "ecs"
+Batch = "batch"
+
+
+class LambdaApi(Enum):
+    Invoke = "invoke"
+
+
+class GlueApi(Enum):
+    StartJobRun = "startJobRun"
+
+
+class EcsApi(Enum):
+    RunTask = "runTask"
+
+
+class BatchApi(Enum):
+    SubmitJob = "submitJob"
 
 
 class LambdaStep(Task):
@@ -44,14 +66,16 @@ def __init__(self, state_id, wait_for_callback=False, **kwargs):
             """
             Example resource arn: arn:aws:states:::lambda:invoke.waitForTaskToken
             """
-            kwargs[Field.Resource.value] =  resource_integration_arn_builder(IntegrationServices.Lambda,
-                                                                             LambdaApi.Invoke,
-                                                                             IntegrationPattern.WaitForTaskToken)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Lambda,
+                                                                        LambdaApi.Invoke,
+                                                                        IntegrationPattern.WaitForTaskToken)
         else:
             """
             Example resource arn: arn:aws:states:::lambda:invoke
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Lambda, LambdaApi.Invoke)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Lambda, LambdaApi.Invoke)
 
 
         super(LambdaStep, self).__init__(state_id, **kwargs)
@@ -82,15 +106,17 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             """
             Example resource arn: arn:aws:states:::glue:startJobRun.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Glue,
-                                                                            GlueApi.StartJobRun,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Glue,
+                                                                       GlueApi.StartJobRun,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::glue:startJobRun
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Glue,
-                                                                            GlueApi.StartJobRun)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Glue,
+                                                                       GlueApi.StartJobRun)
 
         super(GlueStartJobRunStep, self).__init__(state_id, **kwargs)
 
@@ -120,15 +146,17 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             """
             Example resource arn: arn:aws:states:::batch:submitJob.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Batch,
-                                                                            BatchApi.SubmitJob,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Batch,
+                                                                       BatchApi.SubmitJob,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::batch:submitJob
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.Batch,
-                                                                            BatchApi.SubmitJob)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Batch,
+                                                                       BatchApi.SubmitJob)
 
         super(BatchSubmitJobStep, self).__init__(state_id, **kwargs)
 
@@ -158,14 +186,16 @@ def __init__(self, state_id, wait_for_completion=True, **kwargs):
             """
             Example resource arn: arn:aws:states:::ecs:runTask.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.ECS,
-                                                                            EcsApi.RunTask,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Ecs,
+                                                                       EcsApi.RunTask,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::ecs:runTask
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.ECS,
-                                                                            EcsApi.RunTask)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(Ecs,
+                                                                       EcsApi.RunTask)
 
         super(EcsRunTaskStep, self).__init__(state_id, **kwargs)

src/stepfunctions/steps/integration_resources.py

@@ -15,75 +15,16 @@
 
 from enum import Enum
 
-"""
-Enum classes for task integration resource arn builder
-"""
-
 
 class IntegrationPattern(Enum):
+    """
+    Integration pattern enum classes for task integration resource arn builder
+    """
+
     WaitForTaskToken = "waitForTaskToken"
     WaitForCompletion = "sync"
+    RequestResponse = ""
 
 
-class IntegrationServices(Enum):
-    Lambda = "lambda"
-    SageMaker = "sagemaker"
-    Glue = "glue"
-    ECS = "ecs"
-    Batch = "batch"
-    DynamoDB = "dynamodb"
-    SNS = "sns"
-    SQS = "sqs"
-    ElasticMapReduce = "elasticmapreduce"
-
-
-class LambdaApi(Enum):
-    Invoke = "invoke"
-
-
-class SageMakerApi(Enum):
-    CreateTrainingJob = "createTrainingJob"
-    CreateTransformJob = "createTransformJob"
-    CreateModel = "createModel"
-    CreateEndpointConfig = "createEndpointConfig"
-    UpdateEndpoint = "updateEndpoint"
-    CreateEndpoint = "createEndpoint"
-    CreateHyperParameterTuningJob = "createHyperParameterTuningJob"
-    CreateProcessingJob = "createProcessingJob"
-
-
-class GlueApi(Enum):
-    StartJobRun = "startJobRun"
-
-
-class EcsApi(Enum):
-    RunTask = "runTask"
-
-
-class BatchApi(Enum):
-    SubmitJob = "submitJob"
-
-
-class DynamoDBApi(Enum):
-    GetItem = "getItem"
-    PutItem = "putItem"
-    DeleteItem = "deleteItem"
-    UpdateItem = "updateItem"
-
-
-class SnsApi(Enum):
-    Publish = "publish"
-
-
-class SqsApi(Enum):
-    SendMessage = "sendMessage"
 
 
-class ElasticMapReduceApi(Enum):
-    CreateCluster = "createCluster"
-    TerminateCluster = "terminateCluster"
-    AddStep = "addStep"
-    CancelStep = "cancelStep"
-    SetClusterTerminationProtection = "setClusterTerminationProtection"
-    ModifyInstanceFleetByName = "modifyInstanceFleetByName"
-    ModifyInstanceGroupByName = "modifyInstanceGroupByName"

src/stepfunctions/steps/sagemaker.py

@@ -12,16 +12,31 @@
 # permissions and limitations under the License.
 from __future__ import absolute_import
 
+from enum import Enum
 from stepfunctions.inputs import ExecutionInput, StepInput
 from stepfunctions.steps.states import Task
 from stepfunctions.steps.fields import Field
-from stepfunctions.steps.utils import tags_dict_to_kv_list, resource_integration_arn_builder
-from stepfunctions.steps.integration_resources import IntegrationPattern, IntegrationServices, SageMakerApi
+from stepfunctions.steps.utils import tags_dict_to_kv_list, get_service_integration_arn
+from stepfunctions.steps.integration_resources import IntegrationPattern
 
 from sagemaker.workflow.airflow import training_config, transform_config, model_config, tuning_config, processing_config
 from sagemaker.model import Model, FrameworkModel
 from sagemaker.model_monitor import DataCaptureConfig
 
+SageMaker = "sagemaker"
+
+
+class SageMakerApi(Enum):
+    CreateTrainingJob = "createTrainingJob"
+    CreateTransformJob = "createTransformJob"
+    CreateModel = "createModel"
+    CreateEndpointConfig = "createEndpointConfig"
+    UpdateEndpoint = "updateEndpoint"
+    CreateEndpoint = "createEndpoint"
+    CreateHyperParameterTuningJob = "createHyperParameterTuningJob"
+    CreateProcessingJob = "createProcessingJob"
+
+
 class TrainingStep(Task):
 
     """
@@ -62,15 +77,17 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
             """
             Example resource arn: arn:aws:states:::sagemaker:createTrainingJob.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateTrainingJob,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateTrainingJob,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::sagemaker:createTrainingJob
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateTrainingJob)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateTrainingJob)
 
         if isinstance(job_name, str):
             parameters = training_config(estimator=estimator, inputs=data, job_name=job_name, mini_batch_size=mini_batch_size)
@@ -154,15 +171,17 @@ def __init__(self, state_id, transformer, job_name, model_name, data, data_type=
             """
             Example resource arn: arn:aws:states:::sagemaker:createTransformJob.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateTransformJob,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateTransformJob,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::sagemaker:createTransformJob
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateTransformJob)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateTransformJob)
 
         if isinstance(job_name, str):
             parameters = transform_config(
@@ -248,8 +267,9 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
         """
         Example resource arn: arn:aws:states:::sagemaker:createModel
         """
-        kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                        SageMakerApi.CreateModel)
+
+        kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                   SageMakerApi.CreateModel)
 
         super(ModelStep, self).__init__(state_id, **kwargs)
 
@@ -293,8 +313,9 @@ def __init__(self, state_id, endpoint_config_name, model_name, initial_instance_
         """
         Example resource arn: arn:aws:states:::sagemaker:createEndpointConfig
         """
-        kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                        SageMakerApi.CreateEndpointConfig)
+
+        kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                   SageMakerApi.CreateEndpointConfig)
 
         kwargs[Field.Parameters.value] = parameters
 
@@ -330,14 +351,16 @@ def __init__(self, state_id, endpoint_name, endpoint_config_name, tags=None, upd
             """
             Example resource arn: arn:aws:states:::sagemaker:updateEndpoint
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.UpdateEndpoint)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.UpdateEndpoint)
         else:
             """
             Example resource arn: arn:aws:states:::sagemaker:createEndpoint
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateEndpoint)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateEndpoint)
 
         kwargs[Field.Parameters.value] = parameters
 
@@ -378,15 +401,17 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
             """
             Example resource arn: arn:aws:states:::sagemaker:createHyperParameterTuningJob.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateHyperParameterTuningJob,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateHyperParameterTuningJob,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::sagemaker:createHyperParameterTuningJob
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateHyperParameterTuningJob)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateHyperParameterTuningJob)
 
         parameters = tuning_config(tuner=tuner, inputs=data, job_name=job_name).copy()
 
@@ -436,15 +461,17 @@ def __init__(self, state_id, processor, job_name, inputs=None, outputs=None, exp
             """
             Example resource arn: arn:aws:states:::sagemaker:createProcessingJob.sync
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateProcessingJob,
-                                                                            IntegrationPattern.WaitForCompletion)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateProcessingJob,
+                                                                       IntegrationPattern.WaitForCompletion)
         else:
             """
             Example resource arn: arn:aws:states:::sagemaker:createProcessingJob
             """
-            kwargs[Field.Resource.value] = resource_integration_arn_builder(IntegrationServices.SageMaker,
-                                                                            SageMakerApi.CreateProcessingJob)
+
+            kwargs[Field.Resource.value] = get_service_integration_arn(SageMaker,
+                                                                       SageMakerApi.CreateProcessingJob)
 
         if isinstance(job_name, str):
             parameters = processing_config(processor=processor, inputs=inputs, outputs=outputs, container_arguments=container_arguments, container_entrypoint=container_entrypoint, kms_key_id=kms_key_id, job_name=job_name)