Coverage for async-durable-execution/src/async_durable_execution/runner/local/processors/step.py: 100%

1"""Step operation processor for handling STEP operation updates.""" 

2 

3from __future__ import annotations 

4 

5from datetime import datetime, timedelta, timezone 

6from typing import Any 

7 

8from ....models import ( 

9 Operation, 

10 OperationAction, 

11 OperationStatus, 

12 OperationUpdate, 

13 StepDetails, 

14) 

15from .base import ( 

16 OperationProcessor, 

17) 

18from ...exceptions import ( 

19 InvalidParameterValueException, 

20) 

21from ..time_scale import scale_delay 

22 

23VALID_ACTIONS_FOR_STEP = frozenset( 

24 [ 

25 OperationAction.START, 

26 OperationAction.FAIL, 

27 OperationAction.RETRY, 

28 OperationAction.SUCCEED, 

29 ] 

30) 

31 

32 

33class StepProcessor(OperationProcessor): 

34 """Processes STEP operation updates with retry scheduling.""" 

35 

36 valid_actions = VALID_ACTIONS_FOR_STEP 

37 _ALLOWED_STATUS_TO_CLOSE = frozenset( 

38 [ 

39 OperationStatus.STARTED, 

40 OperationStatus.READY, 

41 ] 

42 ) 

43 _ALLOWED_STATUS_TO_START = frozenset( 

44 [ 

45 OperationStatus.READY, 

46 ] 

47 ) 

48 _ALLOWED_STATUS_TO_REATTEMPT = frozenset( 

49 [ 

50 OperationStatus.STARTED, 

51 OperationStatus.READY, 

52 ] 

53 ) 

54 

55 @classmethod 

56 def validate(cls, current_state: Operation | None, update: OperationUpdate) -> None: 

57 """Validate STEP operation update.""" 

58 super().validate(current_state, update) 

59 if current_state is None: 

60 return 

61 

62 match update.action: 

63 case OperationAction.START: 

64 if current_state.status not in cls._ALLOWED_STATUS_TO_START: 

65 msg_step_start: str = "Invalid current STEP state to start." 

66 raise InvalidParameterValueException(msg_step_start) 

67 case OperationAction.FAIL | OperationAction.SUCCEED: 

68 if current_state.status not in cls._ALLOWED_STATUS_TO_CLOSE: 

69 msg_step_close: str = "Invalid current STEP state to close." 

70 raise InvalidParameterValueException(msg_step_close) 

71 if update.action == OperationAction.FAIL and update.payload is not None: 

72 msg_fail_payload: str = "Cannot provide a Payload for FAIL action." 

73 raise InvalidParameterValueException(msg_fail_payload) 

74 if ( 

75 update.action == OperationAction.SUCCEED 

76 and update.error is not None 

77 ): 

78 msg_succeed_error: str = ( 

79 "Cannot provide an Error for SUCCEED action." 

80 ) 

81 raise InvalidParameterValueException(msg_succeed_error) 

82 case OperationAction.RETRY: 

83 if current_state.status not in cls._ALLOWED_STATUS_TO_REATTEMPT: 

84 msg_step_retry: str = "Invalid current STEP state to re-attempt." 

85 raise InvalidParameterValueException(msg_step_retry) 

86 if update.step_options is None: 

87 msg_step_options: str = "Invalid StepOptions for the given action." 

88 raise InvalidParameterValueException(msg_step_options) 

89 if update.error is not None and update.payload is not None: 

90 msg_retry_both: str = ( 

91 "Cannot provide both error and payload to RETRY a STEP." 

92 ) 

93 raise InvalidParameterValueException(msg_retry_both) 

94 

95 def process( 

96 self, 

97 update: OperationUpdate, 

98 current_op: Operation | None, 

99 notifier: Any, 

100 execution_arn: str, 

101 ) -> Operation: 

102 """Process STEP operation update with scheduler integration for retries.""" 

103 match update.action: 

104 case OperationAction.START: 

105 return self._translate_update_to_operation( 

106 update=update, 

107 current_operation=current_op, 

108 status=OperationStatus.STARTED, 

109 ) 

110 case OperationAction.RETRY: 

111 # set Status=PENDING, next attempt time, attempt count + 1 

112 delay = ( 

113 update.step_options.next_attempt_delay_seconds 

114 if update.step_options 

115 else 0 

116 ) 

117 scaled_delay = scale_delay(delay) 

118 next_attempt_time = datetime.now(timezone.utc) + timedelta( 

119 seconds=scaled_delay 

120 ) 

121 

122 # Build new step_details with incremented attempt 

123 current_attempt = ( 

124 current_op.step_details.attempt 

125 if current_op and current_op.step_details 

126 else 0 

127 ) 

128 new_step_details = StepDetails( 

129 attempt=current_attempt + 1, 

130 next_attempt_timestamp=next_attempt_time, 

131 result=( 

132 current_op.step_details.result 

133 if current_op and current_op.step_details 

134 else None 

135 ), 

136 error=( 

137 current_op.step_details.error 

138 if current_op and current_op.step_details 

139 else None 

140 ), 

141 ) 

142 

143 # Create new operation with updated step_details 

144 retry_operation = Operation( 

145 operation_id=update.operation_id, 

146 operation_type=update.operation_type, 

147 status=OperationStatus.PENDING, 

148 parent_id=update.parent_id, 

149 name=update.name, 

150 start_timestamp=( 

151 current_op.start_timestamp 

152 if current_op 

153 else datetime.now(timezone.utc) 

154 ), 

155 end_timestamp=None, 

156 sub_type=update.sub_type, 

157 execution_details=current_op.execution_details 

158 if current_op 

159 else None, 

160 context_details=current_op.context_details if current_op else None, 

161 step_details=new_step_details, 

162 wait_details=current_op.wait_details if current_op else None, 

163 callback_details=current_op.callback_details 

164 if current_op 

165 else None, 

166 chained_invoke_details=current_op.chained_invoke_details 

167 if current_op 

168 else None, 

169 ) 

170 

171 # Schedule step retry timer to fire after delay 

172 notifier.schedule_step_retry( 

173 execution_arn, update.operation_id, scaled_delay 

174 ) 

175 return retry_operation 

176 case OperationAction.SUCCEED: 

177 return self._translate_update_to_operation( 

178 update=update, 

179 current_operation=current_op, 

180 status=OperationStatus.SUCCEEDED, 

181 ) 

182 case OperationAction.FAIL: 

183 return self._translate_update_to_operation( 

184 update=update, 

185 current_operation=current_op, 

186 status=OperationStatus.FAILED, 

187 ) 

188 case _: 

189 msg: str = "Invalid action for STEP operation." 

190 

191 raise InvalidParameterValueException(msg)