From a37b468997660f481243ac2179be18673510d1f8 Mon Sep 17 00:00:00 2001 From: MEPalma Date: Sat, 12 Aug 2023 20:02:52 +0200 Subject: [PATCH] support for sync2 --- .../state_task/service/resource.py | 176 ++++++++------- .../service/state_task_service_callback.py | 7 + .../service/state_task_service_sfn.py | 28 ++- .../templates/callbacks/callback_templates.py | 3 + .../sfn_start_execution_sync2.json5 | 16 ++ .../v2/callback/test_callback.py | 47 ++++ .../v2/callback/test_callback.snapshot.json | 208 ++++++++++++++++++ 7 files changed, 405 insertions(+), 80 deletions(-) create mode 100644 tests/aws/stepfunctions/templates/callbacks/statemachines/sfn_start_execution_sync2.json5 diff --git a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/resource.py b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/resource.py index 2b376412a18b4..f790e8cd03e7c 100644 --- a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/resource.py +++ b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/resource.py @@ -1,7 +1,8 @@ from __future__ import annotations import abc -from typing import Final, Optional, TypedDict +from itertools import takewhile +from typing import Final, Optional from localstack.services.stepfunctions.asl.component.component import Component from localstack.utils.aws import aws_stack @@ -9,86 +10,108 @@ class ResourceCondition(str): WaitForTaskToken = "waitForTaskToken" + Sync2 = "sync:2" Sync = "sync" -class ResourceARN(TypedDict): +class ResourceARN: + arn: str partition: str service: str region: str account: str task_type: str name: str + option: str - -class Resource(Component, abc.ABC): - def __init__(self, resource_arn: str, partition: str, region: str, account: str): - self.resource_arn: Final[str] = resource_arn - self.partition: Final[str] = partition - self.region: Final[str] = region - self.account: Final[str] = account + def __init__( + self, + arn: str, + partition: str, + service: str, + region: str, + account: str, + task_type: str, + name: str, + option: Optional[str], + ): + self.arn = arn + self.partition = partition + self.service = service + self.region = region + self.account = account + self.task_type = task_type + self.name = name + self.option = option @staticmethod - def parse_resource_arn(arn: str) -> ResourceARN: - cmps: list[str] = arn.split(":") - return ResourceARN( - partition=cmps[1], - service=cmps[2], - region=cmps[3], - account=cmps[4], - task_type=cmps[5], - name=cmps[6], + def _consume_until(text: str, symbol: str) -> tuple[str, str]: + value = "".join(takewhile(lambda c: c != symbol, text)) + tail_idx = len(value) + 1 + return value, text[tail_idx:] + + @classmethod + def from_arn(cls, arn: str) -> ResourceARN: + _, arn_tail = ResourceARN._consume_until(arn, ":") + partition, arn_tail = ResourceARN._consume_until(arn_tail, ":") + service, arn_tail = ResourceARN._consume_until(arn_tail, ":") + region, arn_tail = ResourceARN._consume_until(arn_tail, ":") + account, arn_tail = ResourceARN._consume_until(arn_tail, ":") + task_type, arn_tail = ResourceARN._consume_until(arn_tail, ":") + name, arn_tail = ResourceARN._consume_until(arn_tail, ".") + option = arn_tail + return cls( + arn=arn, + partition=partition, + service=service, + region=region, + account=account, + task_type=task_type, + name=name, + option=option, ) + +class Resource(Component, abc.ABC): + resource_arn: Final[str] + partition: Final[str] + region: Final[str] + account: Final[str] + + def __init__(self, resource_arn: ResourceARN): + self.resource_arn = resource_arn.arn + self.partition = resource_arn.partition + self.region = resource_arn.region + self.account = resource_arn.account + @staticmethod def from_resource_arn(arn: str) -> Resource: - resource_arn: ResourceARN = Resource.parse_resource_arn(arn) - if not resource_arn["region"]: - resource_arn["region"] = aws_stack.get_region() - match resource_arn["service"], resource_arn["task_type"]: + resource_arn = ResourceARN.from_arn(arn) + if not resource_arn.region: + resource_arn.region = aws_stack.get_region() + match resource_arn.service, resource_arn.task_type: case "lambda", "function": - return LambdaResource( - resource_arn=arn, - partition=resource_arn["partition"], - region=resource_arn["region"], - account=resource_arn["account"], - function_name=resource_arn["name"], - ) + return LambdaResource(resource_arn=resource_arn) case "states", "activity": - return ActivityResource( - resource_arn=arn, - partition=resource_arn["partition"], - region=resource_arn["region"], - account=resource_arn["account"], - name=resource_arn["name"], - ) - case "states", service_name: - return ServiceResource( - resource_arn=arn, - partition=resource_arn["partition"], - region=resource_arn["region"], - account=resource_arn["account"], - service_name=service_name, # noqa - api_name=resource_arn["name"], - ) + return ActivityResource(resource_arn=resource_arn) + case "states", _: + return ServiceResource(resource_arn=resource_arn) class ActivityResource(Resource): - def __init__(self, resource_arn: str, partition: str, region: str, account: str, name: str): - super().__init__( - resource_arn=resource_arn, partition=partition, region=region, account=account - ) - self.name: str = name + name: Final[str] + + def __init__(self, resource_arn: ResourceARN): + super().__init__(resource_arn=resource_arn) + self.name = resource_arn.name class LambdaResource(Resource): - def __init__( - self, resource_arn: str, partition: str, region: str, account: str, function_name: str - ): - super().__init__( - resource_arn=resource_arn, partition=partition, region=region, account=account - ) - self.function_name: str = function_name + function_name: Final[str] + + def __init__(self, resource_arn: ResourceARN): + super().__init__(resource_arn=resource_arn) + self.function_name: str = resource_arn.name class ServiceResource(Resource): @@ -97,32 +120,29 @@ class ServiceResource(Resource): api_action: Final[str] condition: Final[Optional[str]] - def __init__( - self, - resource_arn: str, - partition: str, - region: str, - account: str, - service_name: str, - api_name: str, - ): - super().__init__( - resource_arn=resource_arn, partition=partition, region=region, account=account - ) - self.service_name = service_name - self.api_name = api_name + def __init__(self, resource_arn: ResourceARN): + super().__init__(resource_arn=resource_arn) + self.service_name = resource_arn.task_type - arn_parts = resource_arn.split(":") - tail_part = arn_parts[-1] - tail_parts = tail_part.split(".") - self.api_action = tail_parts[0] + name_parts = resource_arn.name.split(":") + if len(name_parts) == 1: + self.api_name = self.service_name + self.api_action = resource_arn.name + elif len(name_parts) > 1: + self.api_name = name_parts[0] + self.api_action = name_parts[1] + else: + raise RuntimeError(f"Incorrect definition of ResourceArn.name: '{resource_arn.name}'.") self.condition = None - if len(tail_parts) > 1: - match tail_parts[-1]: - case "waitForTaskToken": + option = resource_arn.option + if option: + match option: + case ResourceCondition.WaitForTaskToken: self.condition = ResourceCondition.WaitForTaskToken case "sync": self.condition = ResourceCondition.Sync + case "sync:2": + self.condition = ResourceCondition.Sync2 case unsupported: raise RuntimeError(f"Unsupported condition '{unsupported}'.") diff --git a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_callback.py b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_callback.py index 0c999a85211d7..4668f906c632d 100644 --- a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_callback.py +++ b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_callback.py @@ -84,6 +84,11 @@ def _sync(self, env: Environment) -> None: f"Unsupported .sync callback procedure in resource {self.resource.resource_arn}" ) + def _sync2(self, env: Environment) -> None: + raise RuntimeError( + f"Unsupported .sync:2 callback procedure in resource {self.resource.resource_arn}" + ) + def _is_condition(self): return self.resource.condition is not None @@ -127,6 +132,8 @@ def _after_eval_execution(self, env: Environment) -> None: self._wait_for_task_token(env=env) case ResourceCondition.Sync: self._sync(env=env) + case ResourceCondition.Sync2: + self._sync2(env=env) case unsupported: raise NotImplementedError(f"Unsupported callback type '{unsupported}'.") diff --git a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_sfn.py b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_sfn.py index 22c0eab82b982..bc93a9f31ed71 100644 --- a/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_sfn.py +++ b/localstack/services/stepfunctions/asl/component/state/state_execution/state_task/service/state_task_service_sfn.py @@ -1,3 +1,4 @@ +import json from typing import Any, Final, Optional from botocore.config import Config @@ -151,6 +152,18 @@ def _normalised_parameters_bindings(self, parameters: dict[str, str]) -> dict[st return normalised_parameters + @staticmethod + def _sync2_api_output_of(typ: type, value: json) -> None: + def _replace_with_json_if_str(key: str) -> None: + inner_value = value.get(key) + if isinstance(inner_value, str): + value[key] = json.loads(inner_value) + + match typ: + case DescribeExecutionOutput: # noqa + _replace_with_json_if_str("input") + _replace_with_json_if_str("output") + def _eval_service_task(self, env: Environment, parameters: dict) -> None: api_action = camel_to_snake_case(self.resource.api_action) sfn_client = self._get_sfn_client() @@ -159,7 +172,7 @@ def _eval_service_task(self, env: Environment, parameters: dict) -> None: self._normalise_botocore_response(self.resource.api_action, response) env.stack.append(response) - def _sync_to_start_machine(self, env: Environment) -> None: + def _sync_to_start_machine(self, env: Environment, sync2_response: bool) -> None: sfn_client = self._get_sfn_client() submission_output: dict = env.stack.pop() @@ -173,6 +186,10 @@ def _has_terminated() -> Optional[dict]: execution_status: ExecutionStatus = describe_execution_output["status"] if execution_status != ExecutionStatus.RUNNING: + if sync2_response: + self._sync2_api_output_of( + typ=DescribeExecutionOutput, value=describe_execution_output + ) self._normalise_botocore_response("describeexecution", describe_execution_output) if execution_status == ExecutionStatus.SUCCEEDED: return describe_execution_output @@ -202,6 +219,13 @@ def _has_terminated() -> Optional[dict]: def _sync(self, env: Environment) -> None: match self.resource.api_action.lower(): case "startexecution": - self._sync_to_start_machine(env=env) + self._sync_to_start_machine(env=env, sync2_response=False) case _: super()._sync(env=env) + + def _sync2(self, env: Environment) -> None: + match self.resource.api_action.lower(): + case "startexecution": + self._sync_to_start_machine(env=env, sync2_response=True) + case _: + super()._sync2(env=env) diff --git a/tests/aws/stepfunctions/templates/callbacks/callback_templates.py b/tests/aws/stepfunctions/templates/callbacks/callback_templates.py index b4be64f6c33de..d3687d4f4df05 100644 --- a/tests/aws/stepfunctions/templates/callbacks/callback_templates.py +++ b/tests/aws/stepfunctions/templates/callbacks/callback_templates.py @@ -10,6 +10,9 @@ class CallbackTemplates(TemplateLoader): SFN_START_EXECUTION_SYNC: Final[str] = os.path.join( _THIS_FOLDER, "statemachines/sfn_start_execution_sync.json5" ) + SFN_START_EXECUTION_SYNC2: Final[str] = os.path.join( + _THIS_FOLDER, "statemachines/sfn_start_execution_sync2.json5" + ) SQS_SUCCESS_ON_TASK_TOKEN: Final[str] = os.path.join( _THIS_FOLDER, "statemachines/sqs_success_on_task_token.json5" ) diff --git a/tests/aws/stepfunctions/templates/callbacks/statemachines/sfn_start_execution_sync2.json5 b/tests/aws/stepfunctions/templates/callbacks/statemachines/sfn_start_execution_sync2.json5 new file mode 100644 index 0000000000000..e28cb6d0b5d5b --- /dev/null +++ b/tests/aws/stepfunctions/templates/callbacks/statemachines/sfn_start_execution_sync2.json5 @@ -0,0 +1,16 @@ +{ + "Comment": "SFN_START_EXECUTION_SYNC:2", + "StartAt": "StartExecution", + "States": { + "StartExecution": { + "Type": "Task", + "Resource": "arn:aws:states:::states:startExecution.sync:2", + "Parameters": { + "StateMachineArn.$": "$.StateMachineArn", + "Input.$": "$.Input", + "Name.$": "$.Name" + }, + "End": true, + } + } +} \ No newline at end of file diff --git a/tests/aws/stepfunctions/v2/callback/test_callback.py b/tests/aws/stepfunctions/v2/callback/test_callback.py index d49da794ca1ae..eed3743ffe7df 100644 --- a/tests/aws/stepfunctions/v2/callback/test_callback.py +++ b/tests/aws/stepfunctions/v2/callback/test_callback.py @@ -235,6 +235,53 @@ def test_start_execution_sync( exec_input, ) + @markers.aws.unknown + def test_start_execution_sync2( + self, + aws_client, + create_iam_role_for_sfn, + create_state_machine, + sfn_snapshot, + ): + sfn_snapshot.add_transformer( + JsonpathTransformer( + jsonpath="$..output.StartDate", + replacement="start-date", + replace_reference=False, + ) + ) + sfn_snapshot.add_transformer( + JsonpathTransformer( + jsonpath="$..output.StopDate", + replacement="stop-date", + replace_reference=False, + ) + ) + + template_target = BT.load_sfn_template(BT.BASE_PASS_RESULT) + definition_target = json.dumps(template_target) + state_machine_arn_target = create( + create_iam_role_for_sfn, + create_state_machine, + sfn_snapshot, + definition_target, + ) + + template = CT.load_sfn_template(CT.SFN_START_EXECUTION_SYNC2) + definition = json.dumps(template) + + exec_input = json.dumps( + {"StateMachineArn": state_machine_arn_target, "Input": None, "Name": "TestStartTarget"} + ) + create_and_record_execution( + aws_client.stepfunctions, + create_iam_role_for_sfn, + create_state_machine, + sfn_snapshot, + definition, + exec_input, + ) + @markers.aws.unknown def test_start_execution_sync_delegate_failure( self, diff --git a/tests/aws/stepfunctions/v2/callback/test_callback.snapshot.json b/tests/aws/stepfunctions/v2/callback/test_callback.snapshot.json index 824533bf04e15..84bc366181fd0 100644 --- a/tests/aws/stepfunctions/v2/callback/test_callback.snapshot.json +++ b/tests/aws/stepfunctions/v2/callback/test_callback.snapshot.json @@ -1115,5 +1115,213 @@ } } } + }, + "tests/aws/stepfunctions/v2/callback/test_callback.py::TestCallback::test_start_execution_sync2": { + "recorded-date": "12-08-2023, 16:45:39", + "recorded-content": { + "get_execution_history": { + "events": [ + { + "executionStartedEventDetails": { + "input": { + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Input": null, + "Name": "TestStartTarget" + }, + "inputDetails": { + "truncated": false + }, + "roleArn": "snf_role_arn" + }, + "id": 1, + "previousEventId": 0, + "timestamp": "timestamp", + "type": "ExecutionStarted" + }, + { + "id": 2, + "previousEventId": 0, + "stateEnteredEventDetails": { + "input": { + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Input": null, + "Name": "TestStartTarget" + }, + "inputDetails": { + "truncated": false + }, + "name": "StartExecution" + }, + "timestamp": "timestamp", + "type": "TaskStateEntered" + }, + { + "id": 3, + "previousEventId": 2, + "taskScheduledEventDetails": { + "parameters": { + "Input": null, + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Name": "TestStartTarget" + }, + "region": "", + "resource": "startExecution.sync:2", + "resourceType": "states" + }, + "timestamp": "timestamp", + "type": "TaskScheduled" + }, + { + "id": 4, + "previousEventId": 3, + "taskStartedEventDetails": { + "resource": "startExecution.sync:2", + "resourceType": "states" + }, + "timestamp": "timestamp", + "type": "TaskStarted" + }, + { + "id": 5, + "previousEventId": 4, + "taskSubmittedEventDetails": { + "output": { + "ExecutionArn": "arn:aws:states::111111111111:execution::TestStartTarget", + "SdkHttpMetadata": { + "AllHttpHeaders": { + "x-amzn-RequestId": [ + "" + ], + "connection": [ + "keep-alive" + ], + "Content-Length": [ + "161" + ], + "Date": "date", + "Content-Type": [ + "application/x-amz-json-1.0" + ] + }, + "HttpHeaders": { + "connection": "keep-alive", + "Content-Length": "161", + "Content-Type": "application/x-amz-json-1.0", + "Date": "date", + "x-amzn-RequestId": "" + }, + "HttpStatusCode": 200 + }, + "SdkResponseMetadata": { + "RequestId": "" + }, + "StartDate": "start-date" + }, + "outputDetails": { + "truncated": false + }, + "resource": "startExecution.sync:2", + "resourceType": "states" + }, + "timestamp": "timestamp", + "type": "TaskSubmitted" + }, + { + "id": 6, + "previousEventId": 5, + "taskSucceededEventDetails": { + "output": { + "ExecutionArn": "arn:aws:states::111111111111:execution::TestStartTarget", + "Input": {}, + "InputDetails": { + "Included": true + }, + "Name": "TestStartTarget", + "Output": { + "Arg1": "argument1" + }, + "OutputDetails": { + "Included": true + }, + "StartDate": "start-date", + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Status": "SUCCEEDED", + "StopDate": "stop-date" + }, + "outputDetails": { + "truncated": false + }, + "resource": "startExecution.sync:2", + "resourceType": "states" + }, + "timestamp": "timestamp", + "type": "TaskSucceeded" + }, + { + "id": 7, + "previousEventId": 6, + "stateExitedEventDetails": { + "name": "StartExecution", + "output": { + "ExecutionArn": "arn:aws:states::111111111111:execution::TestStartTarget", + "Input": {}, + "InputDetails": { + "Included": true + }, + "Name": "TestStartTarget", + "Output": { + "Arg1": "argument1" + }, + "OutputDetails": { + "Included": true + }, + "StartDate": "start-date", + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Status": "SUCCEEDED", + "StopDate": "stop-date" + }, + "outputDetails": { + "truncated": false + } + }, + "timestamp": "timestamp", + "type": "TaskStateExited" + }, + { + "executionSucceededEventDetails": { + "output": { + "ExecutionArn": "arn:aws:states::111111111111:execution::TestStartTarget", + "Input": {}, + "InputDetails": { + "Included": true + }, + "Name": "TestStartTarget", + "Output": { + "Arg1": "argument1" + }, + "OutputDetails": { + "Included": true + }, + "StartDate": "start-date", + "StateMachineArn": "arn:aws:states::111111111111:stateMachine:", + "Status": "SUCCEEDED", + "StopDate": "stop-date" + }, + "outputDetails": { + "truncated": false + } + }, + "id": 8, + "previousEventId": 7, + "timestamp": "timestamp", + "type": "ExecutionSucceeded" + } + ], + "ResponseMetadata": { + "HTTPHeaders": {}, + "HTTPStatusCode": 200 + } + } + } } }