Coverage for async_durable_execution/_runner/local/processors/base.py: 100%

1"""Base processor class for operation transformations.""" 

2 

3from __future__ import annotations 

4 

5import datetime 

6from datetime import timedelta 

7from typing import Any, ClassVar 

8 

9from ...._core import ( 

10 CallbackDetails, 

11 ChainedInvokeDetails, 

12 ContextDetails, 

13 ExecutionDetails, 

14 Operation, 

15 OperationAction, 

16 OperationStatus, 

17 OperationType, 

18 OperationUpdate, 

19 StepDetails, 

20 WaitDetails, 

21) 

22from ...exceptions import ( 

23 InvalidParameterValueException, 

24) 

25 

26 

27class OperationProcessor: 

28 """Base class for processing OperationUpdate to Operation transformations.""" 

29 

30 valid_actions: ClassVar[frozenset[OperationAction]] = frozenset() 

31 

32 @classmethod 

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

34 """Validate operation action and transition before processing.""" 

35 cls._validate_action(update.action) 

36 

37 @classmethod 

38 def _validate_action(cls, action: OperationAction) -> None: 

39 """Validate that an action is supported by this operation processor.""" 

40 if action not in cls.valid_actions: 

41 msg = "Invalid action for the given operation type." 

42 raise InvalidParameterValueException(msg) 

43 

44 def process( 

45 self, 

46 update: OperationUpdate, 

47 current_op: Operation | None, 

48 notifier: Any, 

49 execution_arn: str, 

50 ) -> Operation | None: 

51 """Process an operation update and return the transformed operation.""" 

52 raise NotImplementedError 

53 

54 def _get_start_time( 

55 self, current_operation: Operation | None 

56 ) -> datetime.datetime | None: 

57 start_time: datetime.datetime | None = ( 

58 current_operation.start_timestamp 

59 if current_operation 

60 else datetime.datetime.now(tz=datetime.timezone.utc) 

61 ) 

62 return start_time 

63 

64 def _get_end_time( 

65 self, current_operation: Operation | None, status: OperationStatus 

66 ) -> datetime.datetime | None: 

67 """Get end timestamp for operation based on current state and status.""" 

68 if current_operation and current_operation.end_timestamp: 

69 return current_operation.end_timestamp 

70 if status in { 

71 OperationStatus.SUCCEEDED, 

72 OperationStatus.FAILED, 

73 OperationStatus.CANCELLED, 

74 OperationStatus.TIMED_OUT, 

75 OperationStatus.STOPPED, 

76 }: 

77 return datetime.datetime.now(tz=datetime.timezone.utc) 

78 return None 

79 

80 def _create_execution_details( 

81 self, update: OperationUpdate 

82 ) -> ExecutionDetails | None: 

83 """Create ExecutionDetails from OperationUpdate.""" 

84 return ( 

85 ExecutionDetails(input_payload=update.payload) 

86 if update.operation_type == OperationType.EXECUTION 

87 else None 

88 ) 

89 

90 def _create_context_details(self, update: OperationUpdate) -> ContextDetails | None: 

91 """Create ContextDetails from OperationUpdate.""" 

92 return ( 

93 ContextDetails( 

94 result=update.payload, 

95 error=update.error, 

96 replay_children=update.context_options.replay_children 

97 if update.context_options 

98 else False, 

99 ) 

100 if update.operation_type == OperationType.CONTEXT 

101 else None 

102 ) 

103 

104 def _create_step_details( 

105 self, 

106 update: OperationUpdate, 

107 current_operation: Operation | None = None, 

108 ) -> StepDetails | None: 

109 """Create StepDetails from OperationUpdate. 

110 

111 Automatically increments attempt count for RETRY, SUCCEED, and FAIL actions. 

112 """ 

113 

114 attempt: int = 0 

115 next_attempt_timestamp: datetime.datetime | None = None 

116 

117 if update.operation_type is OperationType.STEP: 

118 if current_operation and current_operation.step_details: 

119 attempt = current_operation.step_details.attempt 

120 next_attempt_timestamp = ( 

121 current_operation.step_details.next_attempt_timestamp 

122 ) 

123 # Increment attempt for RETRY, SUCCEED, and FAIL actions 

124 if update.action in { 

125 OperationAction.RETRY, 

126 OperationAction.SUCCEED, 

127 OperationAction.FAIL, 

128 }: 

129 attempt += 1 

130 result = update.payload 

131 error = update.error 

132 if ( 

133 update.action is OperationAction.START 

134 and current_operation is not None 

135 and current_operation.step_details is not None 

136 ): 

137 if result is None: 

138 result = current_operation.step_details.result 

139 return StepDetails( 

140 attempt=attempt, 

141 next_attempt_timestamp=next_attempt_timestamp, 

142 result=result, 

143 error=error, 

144 ) 

145 

146 return None 

147 

148 def _create_callback_details( 

149 self, update: OperationUpdate 

150 ) -> CallbackDetails | None: 

151 """Create CallbackDetails from OperationUpdate.""" 

152 return ( 

153 CallbackDetails( 

154 callback_id="placeholder", result=update.payload, error=update.error 

155 ) 

156 if update.operation_type == OperationType.CALLBACK 

157 else None 

158 ) 

159 

160 def _create_invoke_details( 

161 self, update: OperationUpdate 

162 ) -> ChainedInvokeDetails | None: 

163 """Create ChainedInvokeDetails from OperationUpdate.""" 

164 if ( 

165 update.operation_type == OperationType.CHAINED_INVOKE 

166 and update.chained_invoke_options 

167 ): 

168 return ChainedInvokeDetails(result=update.payload, error=update.error) 

169 return None 

170 

171 def _translate_update_to_operation( 

172 self, 

173 update: OperationUpdate, 

174 current_operation: Operation | None, 

175 status: OperationStatus, 

176 ) -> Operation: 

177 """Transform OperationUpdate to Operation, always creating new Operation.""" 

178 start_time: datetime.datetime | None = self._get_start_time(current_operation) 

179 end_time: datetime.datetime | None = self._get_end_time( 

180 current_operation, status 

181 ) 

182 

183 execution_details = self._create_execution_details(update) 

184 context_details = self._create_context_details(update) 

185 step_details = self._create_step_details(update, current_operation) 

186 callback_details = self._create_callback_details(update) 

187 invoke_details = self._create_invoke_details(update) 

188 wait_details = self._create_wait_details(update, current_operation) 

189 

190 return Operation( 

191 operation_id=update.operation_id, 

192 parent_id=update.parent_id, 

193 name=update.name, 

194 start_timestamp=start_time, 

195 end_timestamp=end_time, 

196 operation_type=update.operation_type, 

197 status=status, 

198 sub_type=update.sub_type, 

199 execution_details=execution_details, 

200 context_details=context_details, 

201 step_details=step_details, 

202 callback_details=callback_details, 

203 chained_invoke_details=invoke_details, 

204 wait_details=wait_details, 

205 ) 

206 

207 def _create_wait_details( 

208 self, update: OperationUpdate, current_operation: Operation | None 

209 ) -> WaitDetails | None: 

210 """Create WaitDetails from OperationUpdate.""" 

211 if update.operation_type == OperationType.WAIT and update.wait_options: 

212 if current_operation and current_operation.wait_details: 

213 scheduled_end_timestamp = ( 

214 current_operation.wait_details.scheduled_end_timestamp 

215 ) 

216 else: 

217 scheduled_end_timestamp = datetime.datetime.now( 

218 tz=datetime.timezone.utc 

219 ) + timedelta(seconds=update.wait_options.wait_seconds) 

220 return WaitDetails(scheduled_end_timestamp=scheduled_end_timestamp) 

221 return None