1afae7ad-fa95-4b95-b399-21b68a0908b4

Committed 24 Jan 2025 06:19PM UTC coverage: 86.904% (+0.02%) from 86.884%

Build # 1afae7ad-fa95-4b95-b399-21b68a0908b4

Build Type

push

circleci

Committed by

web-flow

Commit Message

StepFunctions: Support for Output Blocks in Choice Rules, Improvments to JSONata Choice Defaults (#12075)

Coverage Stats

44 of 47 new or added lines in 4 files covered. (93.62%)

13 existing lines in 6 files now uncovered.

61132 of 70344 relevant lines covered (86.9%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.15

/localstack-core/localstack/services/stepfunctions/asl/component/state/state.py

from __future__ import annotations

import abc
import datetime
import json
import logging
from abc import ABC
from typing import Final, Optional, Union

from localstack.aws.api.stepfunctions import (
    ExecutionFailedEventDetails,
    HistoryEventExecutionDataDetails,
    HistoryEventType,
    StateEnteredEventDetails,
    StateExitedEventDetails,
    TaskFailedEventDetails,
)
from localstack.services.stepfunctions.asl.component.common.assign.assign_decl import AssignDecl
from localstack.services.stepfunctions.asl.component.common.comment import Comment
from localstack.services.stepfunctions.asl.component.common.error_name.failure_event import (
    FailureEvent,
    FailureEventException,
)
from localstack.services.stepfunctions.asl.component.common.error_name.states_error_name import (
    StatesErrorName,
)
from localstack.services.stepfunctions.asl.component.common.error_name.states_error_name_type import (
    StatesErrorNameType,
)
from localstack.services.stepfunctions.asl.component.common.flow.end import End
from localstack.services.stepfunctions.asl.component.common.flow.next import Next
from localstack.services.stepfunctions.asl.component.common.outputdecl import Output
from localstack.services.stepfunctions.asl.component.common.path.input_path import (
    InputPath,
)
from localstack.services.stepfunctions.asl.component.common.path.output_path import OutputPath
from localstack.services.stepfunctions.asl.component.common.query_language import (
    QueryLanguage,
    QueryLanguageMode,
)
from localstack.services.stepfunctions.asl.component.common.string.string_expression import (
    JSONPATH_ROOT_PATH,
    StringJsonPath,
)
from localstack.services.stepfunctions.asl.component.eval_component import EvalComponent
from localstack.services.stepfunctions.asl.component.state.state_continue_with import (
    ContinueWith,
    ContinueWithEnd,
    ContinueWithNext,
)
from localstack.services.stepfunctions.asl.component.state.state_props import StateProps
from localstack.services.stepfunctions.asl.component.state.state_type import StateType
from localstack.services.stepfunctions.asl.eval.environment import Environment
from localstack.services.stepfunctions.asl.eval.event.event_detail import EventDetails
from localstack.services.stepfunctions.asl.eval.program_state import ProgramRunning
from localstack.services.stepfunctions.asl.eval.states import StateData
from localstack.services.stepfunctions.asl.utils.encoding import to_json_str
from localstack.services.stepfunctions.asl.utils.json_path import NoSuchJsonPathError
from localstack.services.stepfunctions.quotas import is_within_size_quota

LOG = logging.getLogger(__name__)


class CommonStateField(EvalComponent, ABC):
    name: str

    query_language: QueryLanguage

    # The state's type.
    state_type: StateType

    # There can be any number of terminal states per state machine. Only one of Next or End can
    # be used in a state. Some state types, such as Choice, don't support or use the End field.
    continue_with: ContinueWith

    # Holds a human-readable description of the state.
    comment: Optional[Comment]

    # A path that selects a portion of the state's input to be passed to the state's state_task for processing.
    # If omitted, it has the value $ which designates the entire input.
    input_path: Optional[InputPath]

    # A path that selects a portion of the state's output to be passed to the next state.
    # If omitted, it has the value $ which designates the entire output.
    output_path: Optional[OutputPath]

    assign_decl: Optional[AssignDecl]

    output: Optional[Output]

    state_entered_event_type: Final[HistoryEventType]
    state_exited_event_type: Final[Optional[HistoryEventType]]

    def __init__(
        self,
        state_entered_event_type: HistoryEventType,
        state_exited_event_type: Optional[HistoryEventType],
    ):
        self.state_entered_event_type = state_entered_event_type
        self.state_exited_event_type = state_exited_event_type

    def from_state_props(self, state_props: StateProps) -> None:
        self.name = state_props.name
        self.query_language = state_props.get(QueryLanguage) or QueryLanguage()
        self.state_type = state_props.get(StateType)
        self.continue_with = (
            ContinueWithEnd() if state_props.get(End) else ContinueWithNext(state_props.get(Next))
        )
        self.comment = state_props.get(Comment)
        self.assign_decl = state_props.get(AssignDecl)
        # JSONPath sub-productions.
        if self.query_language.query_language_mode == QueryLanguageMode.JSONPath:
            self.input_path = state_props.get(InputPath) or InputPath(
                StringJsonPath(JSONPATH_ROOT_PATH)
            )
            self.output_path = state_props.get(OutputPath) or OutputPath(
                StringJsonPath(JSONPATH_ROOT_PATH)
            )
            self.output = None
        # JSONata sub-productions.
        else:
            self.input_path = None
            self.output_path = None
            self.output = state_props.get(Output)

    def _set_next(self, env: Environment) -> None:
        if env.next_state_name != self.name:
            # Next was already overridden.
            return

        if isinstance(self.continue_with, ContinueWithNext):
            env.next_state_name = self.continue_with.next_state.name
        elif isinstance(self.continue_with, ContinueWithEnd):  # This includes ContinueWithSuccess
            env.set_ended()
        else:
            LOG.error("Could not handle ContinueWith type of '%s'.", type(self.continue_with))

    def _is_language_query_jsonpath(self) -> bool:
        return self.query_language.query_language_mode == QueryLanguageMode.JSONPath

    def _get_state_entered_event_details(self, env: Environment) -> StateEnteredEventDetails:
        return StateEnteredEventDetails(
            name=self.name,
            input=to_json_str(env.states.get_input(), separators=(",", ":")),
            inputDetails=HistoryEventExecutionDataDetails(
                truncated=False  # Always False for api calls.
            ),
        )

    def _get_state_exited_event_details(self, env: Environment) -> StateExitedEventDetails:
        event_details = StateExitedEventDetails(
            name=self.name,
            output=to_json_str(env.states.get_input(), separators=(",", ":")),
            outputDetails=HistoryEventExecutionDataDetails(
                truncated=False  # Always False for api calls.
            ),
        )
        # TODO add typing when these become available in boto.
        assigned_variables = env.variable_store.get_assigned_variables()
        env.variable_store.reset_tracing()
        if assigned_variables:
            event_details["assignedVariables"] = assigned_variables  # noqa
            event_details["assignedVariablesDetails"] = {"truncated": False}  # noqa
        return event_details

    def _verify_size_quota(self, env: Environment, value: Union[str, json]) -> None:
        is_within: bool = is_within_size_quota(value)
        if is_within:
            return
        error_type = StatesErrorNameType.StatesStatesDataLimitExceeded
        cause = (
            f"The state/task '{self.name}' returned a result with a size exceeding "
            f"the maximum number of bytes service limit."
        )
        raise FailureEventException(
            failure_event=FailureEvent(
                env=env,
                error_name=StatesErrorName(typ=error_type),
                event_type=HistoryEventType.TaskFailed,
                event_details=EventDetails(
                    executionFailedEventDetails=ExecutionFailedEventDetails(
                        error=error_type.to_name(),
                        cause=cause,
                    )
                ),
            )
        )

    def _eval_state_input(self, env: Environment) -> None:
        # Filter the input onto the stack.
        if self.input_path:
            self.input_path.eval(env)
        else:
            env.stack.append(env.states.get_input())

    @abc.abstractmethod
    def _eval_state(self, env: Environment) -> None: ...

    def _eval_state_output(self, env: Environment) -> None:
        # Process output value as next state input.
        if self.output_path:
            self.output_path.eval(env=env)
        elif self.output:
            self.output.eval(env=env)
        else:
            current_output = env.stack.pop()
            env.states.reset(input_value=current_output)

    def _eval_body(self, env: Environment) -> None:
        env.event_manager.add_event(
            context=env.event_history_context,
            event_type=self.state_entered_event_type,
            event_details=EventDetails(
                stateEnteredEventDetails=self._get_state_entered_event_details(env=env)
            ),
        )
        env.states.context_object.context_object_data["State"] = StateData(
            EnteredTime=datetime.datetime.now(tz=datetime.timezone.utc).isoformat(), Name=self.name
        )

        self._eval_state_input(env=env)

        try:
            self._eval_state(env)
        except NoSuchJsonPathError as no_such_json_path_error:
            data_json_str = to_json_str(no_such_json_path_error.data)
            cause = (
                f"The JSONPath '{no_such_json_path_error.json_path}' specified for the field '{env.next_field_name}' "
                f"could not be found in the input '{data_json_str}'"
            )
            raise FailureEventException(
                failure_event=FailureEvent(
                    env=env,
                    error_name=StatesErrorName(typ=StatesErrorNameType.StatesRuntime),
                    event_type=HistoryEventType.TaskFailed,
                    event_details=EventDetails(
                        taskFailedEventDetails=TaskFailedEventDetails(
                            error=StatesErrorNameType.StatesRuntime.to_name(), cause=cause
                        )
                    ),
                )
            )

        if not isinstance(env.program_state(), ProgramRunning):
            return

        self._eval_state_output(env=env)

        self._verify_size_quota(env=env, value=env.states.get_input())

        self._set_next(env)

        if self.state_exited_event_type is not None:
            env.event_manager.add_event(
                context=env.event_history_context,
                event_type=self.state_exited_event_type,
                event_details=EventDetails(
                    stateExitedEventDetails=self._get_state_exited_event_details(env=env),
                ),
            )

1	from __future__ import annotations	1✔
2
3	import abc	1✔
4	import datetime	1✔
5	import json	1✔
6	import logging	1✔
7	from abc import ABC	1✔
8	from typing import Final, Optional, Union	1✔
9
10	from localstack.aws.api.stepfunctions import (	1✔
11	ExecutionFailedEventDetails,
12	HistoryEventExecutionDataDetails,
13	HistoryEventType,
14	StateEnteredEventDetails,
15	StateExitedEventDetails,
16	TaskFailedEventDetails,
17	)
18	from localstack.services.stepfunctions.asl.component.common.assign.assign_decl import AssignDecl	1✔
19	from localstack.services.stepfunctions.asl.component.common.comment import Comment	1✔
20	from localstack.services.stepfunctions.asl.component.common.error_name.failure_event import (	1✔
21	FailureEvent,
22	FailureEventException,
23	)
24	from localstack.services.stepfunctions.asl.component.common.error_name.states_error_name import (	1✔
25	StatesErrorName,
26	)
27	from localstack.services.stepfunctions.asl.component.common.error_name.states_error_name_type import (	1✔
28	StatesErrorNameType,
29	)
30	from localstack.services.stepfunctions.asl.component.common.flow.end import End	1✔
31	from localstack.services.stepfunctions.asl.component.common.flow.next import Next	1✔
32	from localstack.services.stepfunctions.asl.component.common.outputdecl import Output	1✔
33	from localstack.services.stepfunctions.asl.component.common.path.input_path import (	1✔
34	InputPath,
35	)
36	from localstack.services.stepfunctions.asl.component.common.path.output_path import OutputPath	1✔
37	from localstack.services.stepfunctions.asl.component.common.query_language import (	1✔
38	QueryLanguage,
39	QueryLanguageMode,
40	)
41	from localstack.services.stepfunctions.asl.component.common.string.string_expression import (	1✔
42	JSONPATH_ROOT_PATH,
43	StringJsonPath,
44	)
45	from localstack.services.stepfunctions.asl.component.eval_component import EvalComponent	1✔
46	from localstack.services.stepfunctions.asl.component.state.state_continue_with import (	1✔
47	ContinueWith,
48	ContinueWithEnd,
49	ContinueWithNext,
50	)
51	from localstack.services.stepfunctions.asl.component.state.state_props import StateProps	1✔
52	from localstack.services.stepfunctions.asl.component.state.state_type import StateType	1✔
53	from localstack.services.stepfunctions.asl.eval.environment import Environment	1✔
54	from localstack.services.stepfunctions.asl.eval.event.event_detail import EventDetails	1✔
55	from localstack.services.stepfunctions.asl.eval.program_state import ProgramRunning	1✔
56	from localstack.services.stepfunctions.asl.eval.states import StateData	1✔
57	from localstack.services.stepfunctions.asl.utils.encoding import to_json_str	1✔
58	from localstack.services.stepfunctions.asl.utils.json_path import NoSuchJsonPathError	1✔
59	from localstack.services.stepfunctions.quotas import is_within_size_quota	1✔
60
61	LOG = logging.getLogger(__name__)	1✔
62
63
64	class CommonStateField(EvalComponent, ABC):	1✔
65	name: str	1✔
66
67	query_language: QueryLanguage	1✔
68
69	# The state's type.
70	state_type: StateType	1✔
71
72	# There can be any number of terminal states per state machine. Only one of Next or End can
73	# be used in a state. Some state types, such as Choice, don't support or use the End field.
74	continue_with: ContinueWith	1✔
75
76	# Holds a human-readable description of the state.
77	comment: Optional[Comment]	1✔
78
79	# A path that selects a portion of the state's input to be passed to the state's state_task for processing.
80	# If omitted, it has the value $ which designates the entire input.
81	input_path: Optional[InputPath]	1✔
82
83	# A path that selects a portion of the state's output to be passed to the next state.
84	# If omitted, it has the value $ which designates the entire output.
85	output_path: Optional[OutputPath]	1✔
86
87	assign_decl: Optional[AssignDecl]	1✔
88
89	output: Optional[Output]	1✔
90
91	state_entered_event_type: Final[HistoryEventType]	1✔
92	state_exited_event_type: Final[Optional[HistoryEventType]]	1✔
93
94	def __init__(	1✔
95	self,
96	state_entered_event_type: HistoryEventType,
97	state_exited_event_type: Optional[HistoryEventType],
98	):
99	self.state_entered_event_type = state_entered_event_type	1✔
100	self.state_exited_event_type = state_exited_event_type	1✔
101
102	def from_state_props(self, state_props: StateProps) -> None:	1✔
103	self.name = state_props.name	1✔
104	self.query_language = state_props.get(QueryLanguage) or QueryLanguage()	1✔
105	self.state_type = state_props.get(StateType)	1✔
106	self.continue_with = (	1✔
107	ContinueWithEnd() if state_props.get(End) else ContinueWithNext(state_props.get(Next))
108	)
109	self.comment = state_props.get(Comment)	1✔
110	self.assign_decl = state_props.get(AssignDecl)	1✔
111	# JSONPath sub-productions.
112	if self.query_language.query_language_mode == QueryLanguageMode.JSONPath:	1✔
113	self.input_path = state_props.get(InputPath) or InputPath(	1✔
114	StringJsonPath(JSONPATH_ROOT_PATH)
115	)
116	self.output_path = state_props.get(OutputPath) or OutputPath(	1✔
117	StringJsonPath(JSONPATH_ROOT_PATH)
118	)
119	self.output = None	1✔
120	# JSONata sub-productions.
121	else:
122	self.input_path = None	1✔
123	self.output_path = None	1✔
124	self.output = state_props.get(Output)	1✔
125
126	def _set_next(self, env: Environment) -> None:	1✔
127	if env.next_state_name != self.name:	1✔
128	# Next was already overridden.
129	return	1✔
130
131	if isinstance(self.continue_with, ContinueWithNext):	1✔
132	env.next_state_name = self.continue_with.next_state.name	1✔
133	elif isinstance(self.continue_with, ContinueWithEnd): # This includes ContinueWithSuccess	1✔
134	env.set_ended()	1✔
135	else:
UNCOV 136	LOG.error("Could not handle ContinueWith type of '%s'.", type(self.continue_with))	×
137
138	def _is_language_query_jsonpath(self) -> bool:	1✔
139	return self.query_language.query_language_mode == QueryLanguageMode.JSONPath	1✔
140
141	def _get_state_entered_event_details(self, env: Environment) -> StateEnteredEventDetails:	1✔
142	return StateEnteredEventDetails(	1✔
143	name=self.name,
144	input=to_json_str(env.states.get_input(), separators=(",", ":")),
145	inputDetails=HistoryEventExecutionDataDetails(
146	truncated=False # Always False for api calls.
147	),
148	)
149
150	def _get_state_exited_event_details(self, env: Environment) -> StateExitedEventDetails:	1✔
151	event_details = StateExitedEventDetails(	1✔
152	name=self.name,
153	output=to_json_str(env.states.get_input(), separators=(",", ":")),
154	outputDetails=HistoryEventExecutionDataDetails(
155	truncated=False # Always False for api calls.
156	),
157	)
158	# TODO add typing when these become available in boto.
159	assigned_variables = env.variable_store.get_assigned_variables()	1✔
160	env.variable_store.reset_tracing()	1✔
161	if assigned_variables:	1✔
162	event_details["assignedVariables"] = assigned_variables # noqa	1✔
163	event_details["assignedVariablesDetails"] = {"truncated": False} # noqa	1✔
164	return event_details	1✔
165
166	def _verify_size_quota(self, env: Environment, value: Union[str, json]) -> None:	1✔
167	is_within: bool = is_within_size_quota(value)	1✔
168	if is_within:	1✔
169	return	1✔
170	error_type = StatesErrorNameType.StatesStatesDataLimitExceeded	1✔
171	cause = (	1✔
172	f"The state/task '{self.name}' returned a result with a size exceeding "
173	f"the maximum number of bytes service limit."
174	)
175	raise FailureEventException(	1✔
176	failure_event=FailureEvent(
177	env=env,
178	error_name=StatesErrorName(typ=error_type),
179	event_type=HistoryEventType.TaskFailed,
180	event_details=EventDetails(
181	executionFailedEventDetails=ExecutionFailedEventDetails(
182	error=error_type.to_name(),
183	cause=cause,
184	)
185	),
186	)
187	)
188
189	def _eval_state_input(self, env: Environment) -> None:	1✔
190	# Filter the input onto the stack.
191	if self.input_path:	1✔
192	self.input_path.eval(env)	1✔
193	else:
194	env.stack.append(env.states.get_input())	1✔
195
196	@abc.abstractmethod	1✔
197	def _eval_state(self, env: Environment) -> None: ...	1✔
198
199	def _eval_state_output(self, env: Environment) -> None:	1✔
200	# Process output value as next state input.
201	if self.output_path:	1✔
202	self.output_path.eval(env=env)	1✔
203	elif self.output:	1✔
204	self.output.eval(env=env)	1✔
205	else:
206	current_output = env.stack.pop()	1✔
207	env.states.reset(input_value=current_output)	1✔
208
209	def _eval_body(self, env: Environment) -> None:	1✔
210	env.event_manager.add_event(	1✔
211	context=env.event_history_context,
212	event_type=self.state_entered_event_type,
213	event_details=EventDetails(
214	stateEnteredEventDetails=self._get_state_entered_event_details(env=env)
215	),
216	)
217	env.states.context_object.context_object_data["State"] = StateData(	1✔
218	EnteredTime=datetime.datetime.now(tz=datetime.timezone.utc).isoformat(), Name=self.name
219	)
220
221	self._eval_state_input(env=env)	1✔
222
223	try:	1✔
224	self._eval_state(env)	1✔
225	except NoSuchJsonPathError as no_such_json_path_error:	1✔
226	data_json_str = to_json_str(no_such_json_path_error.data)	1✔
227	cause = (	1✔
228	f"The JSONPath '{no_such_json_path_error.json_path}' specified for the field '{env.next_field_name}' "
229	f"could not be found in the input '{data_json_str}'"
230	)
231	raise FailureEventException(	1✔
232	failure_event=FailureEvent(
233	env=env,
234	error_name=StatesErrorName(typ=StatesErrorNameType.StatesRuntime),
235	event_type=HistoryEventType.TaskFailed,
236	event_details=EventDetails(
237	taskFailedEventDetails=TaskFailedEventDetails(
238	error=StatesErrorNameType.StatesRuntime.to_name(), cause=cause
239	)
240	),
241	)
242	)
243
244	if not isinstance(env.program_state(), ProgramRunning):	1✔
245	return	1✔
246
247	self._eval_state_output(env=env)	1✔
248
249	self._verify_size_quota(env=env, value=env.states.get_input())	1✔
250
251	self._set_next(env)	1✔
252
253	if self.state_exited_event_type is not None:	1✔
254	env.event_manager.add_event(	1✔
255	context=env.event_history_context,
256	event_type=self.state_exited_event_type,
257	event_details=EventDetails(
258	stateExitedEventDetails=self._get_state_exited_event_details(env=env),
259	),
260	)

localstack / localstack / 1afae7ad-fa95-4b95-b399-21b68a0908b4

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous