Coverage for async-durable-execution/src/async_durable_execution/runner/local/model.py: 99%

1"""Models used only by the local durable execution runner.""" 

2 

3from __future__ import annotations 

4 

5import base64 

6import json 

7from dataclasses import dataclass 

8from typing import Any, Protocol 

9 

10from ...execution import DurableExecutionInvocationInput 

11from ...models import ( 

12 CheckpointUpdatedExecutionState, 

13 LambdaContext as LambdaContextProtocol, 

14 Operation, 

15) 

16from ..exceptions import InvalidParameterValueException 

17from ..model import InvokeResponse 

18 

19 

20@dataclass(frozen=True) 

21class LambdaContext(LambdaContextProtocol): 

22 """Lambda context for local testing.""" 

23 

24 aws_request_id: str 

25 log_group_name: str | None = None 

26 log_stream_name: str | None = None 

27 function_name: str | None = None 

28 memory_limit_in_mb: str | None = None 

29 function_version: str | None = None 

30 invoked_function_arn: str | None = None 

31 tenant_id: str | None = None 

32 client_context: dict | None = None 

33 identity: dict | None = None 

34 

35 def get_remaining_time_in_millis(self) -> int: 

36 return 900000 

37 

38 def log(self, msg) -> None: 

39 pass 

40 

41 

42@dataclass(frozen=True) 

43class StartDurableExecutionInput: 

44 """Input for starting a local durable execution.""" 

45 

46 account_id: str 

47 function_name: str 

48 function_qualifier: str 

49 execution_name: str 

50 execution_timeout_seconds: int 

51 execution_retention_period_days: int 

52 invocation_id: str | None = None 

53 trace_fields: dict | None = None 

54 tenant_id: str | None = None 

55 input: str | None = None 

56 lambda_endpoint: str | None = None 

57 

58 @classmethod 

59 def from_dict(cls, data: dict) -> StartDurableExecutionInput: 

60 required_fields = [ 

61 "AccountId", 

62 "FunctionName", 

63 "FunctionQualifier", 

64 "ExecutionName", 

65 "ExecutionTimeoutSeconds", 

66 "ExecutionRetentionPeriodDays", 

67 ] 

68 

69 for field in required_fields: 

70 if field not in data: 

71 msg = f"Missing required field: {field}" 

72 raise InvalidParameterValueException(msg) 

73 

74 return cls( 

75 account_id=data["AccountId"], 

76 function_name=data["FunctionName"], 

77 function_qualifier=data["FunctionQualifier"], 

78 execution_name=data["ExecutionName"], 

79 execution_timeout_seconds=data["ExecutionTimeoutSeconds"], 

80 execution_retention_period_days=data["ExecutionRetentionPeriodDays"], 

81 invocation_id=data.get("InvocationId"), 

82 trace_fields=data.get("TraceFields"), 

83 tenant_id=data.get("TenantId"), 

84 input=data.get("Input"), 

85 lambda_endpoint=data.get("LambdaEndpoint"), 

86 ) 

87 

88 def to_dict(self) -> dict[str, Any]: 

89 result = { 

90 "AccountId": self.account_id, 

91 "FunctionName": self.function_name, 

92 "FunctionQualifier": self.function_qualifier, 

93 "ExecutionName": self.execution_name, 

94 "ExecutionTimeoutSeconds": self.execution_timeout_seconds, 

95 "ExecutionRetentionPeriodDays": self.execution_retention_period_days, 

96 } 

97 if self.invocation_id is not None: 

98 result["InvocationId"] = self.invocation_id 

99 if self.trace_fields is not None: 

100 result["TraceFields"] = self.trace_fields 

101 if self.tenant_id is not None: 

102 result["TenantId"] = self.tenant_id 

103 if self.input is not None: 

104 result["Input"] = self.input 

105 if self.lambda_endpoint is not None: 

106 result["LambdaEndpoint"] = self.lambda_endpoint 

107 return result 

108 

109 def get_normalized_input(self): 

110 """Normalize input string to be JSON deserializable.""" 

111 try: 

112 json.loads(self.input) 

113 return self.input 

114 except (json.JSONDecodeError, TypeError): 

115 return json.dumps(self.input) 

116 

117 

118@dataclass(frozen=True) 

119class StartDurableExecutionOutput: 

120 """Output from starting a local durable execution.""" 

121 

122 execution_arn: str | None = None 

123 

124 @classmethod 

125 def from_dict(cls, data: dict) -> StartDurableExecutionOutput: 

126 return cls(execution_arn=data.get("ExecutionArn")) 

127 

128 def to_dict(self) -> dict[str, Any]: 

129 result = {} 

130 if self.execution_arn is not None: 

131 result["ExecutionArn"] = self.execution_arn 

132 return result 

133 

134 

135@dataclass(frozen=True) 

136class GetDurableExecutionStateResponse: 

137 """Local response containing durable execution state operations.""" 

138 

139 operations: list[Operation] 

140 next_marker: str | None = None 

141 

142 @classmethod 

143 def from_dict(cls, data: dict) -> GetDurableExecutionStateResponse: 

144 operations = [ 

145 Operation.from_dict(op_data) for op_data in data.get("Operations", []) 

146 ] 

147 return cls( 

148 operations=operations, 

149 next_marker=data.get("NextMarker"), 

150 ) 

151 

152 def to_dict(self) -> dict[str, Any]: 

153 result: dict[str, Any] = { 

154 "Operations": [op.to_dict() for op in self.operations] 

155 } 

156 if self.next_marker is not None: 

157 result["NextMarker"] = self.next_marker 

158 return result 

159 

160 

161@dataclass(frozen=True) 

162class SendDurableExecutionCallbackSuccessResponse: 

163 """Response from sending local callback success.""" 

164 

165 

166@dataclass(frozen=True) 

167class SendDurableExecutionCallbackFailureResponse: 

168 """Response from sending local callback failure.""" 

169 

170 

171@dataclass(frozen=True) 

172class SendDurableExecutionCallbackHeartbeatResponse: 

173 """Response from sending local callback heartbeat.""" 

174 

175 

176@dataclass(frozen=True) 

177class CheckpointDurableExecutionResponse: 

178 """Local response from checkpointing a durable execution.""" 

179 

180 checkpoint_token: str 

181 new_execution_state: CheckpointUpdatedExecutionState | None = None 

182 

183 @classmethod 

184 def from_dict(cls, data: dict) -> CheckpointDurableExecutionResponse: 

185 new_execution_state = None 

186 if state_data := data.get("NewExecutionState"): 

187 new_execution_state = CheckpointUpdatedExecutionState.from_dict(state_data) 

188 

189 return cls( 

190 checkpoint_token=data["CheckpointToken"], 

191 new_execution_state=new_execution_state, 

192 ) 

193 

194 def to_dict(self) -> dict[str, Any]: 

195 result: dict[str, Any] = {"CheckpointToken": self.checkpoint_token} 

196 if self.new_execution_state is not None: 

197 result["NewExecutionState"] = self.new_execution_state.to_dict() 

198 return result 

199 

200 

201class Invoker(Protocol): 

202 def create_invocation_input( 

203 self, 

204 *, 

205 start_input: StartDurableExecutionInput, 

206 durable_execution_arn: str, 

207 checkpoint_token: str, 

208 operations: list[Operation], 

209 ) -> DurableExecutionInvocationInput: ... # pragma: no cover 

210 

211 async def invoke( 

212 self, 

213 function_name: str, 

214 input: DurableExecutionInvocationInput, 

215 endpoint_url: str | None = None, 

216 ) -> InvokeResponse: ... # pragma: no cover 

217 

218 def update_endpoint( 

219 self, endpoint_url: str, region_name: str 

220 ) -> None: ... # pragma: no cover 

221 

222 

223@dataclass(frozen=True) 

224class CheckpointToken: 

225 """Model a local checkpoint token.""" 

226 

227 execution_arn: str 

228 token_sequence: int 

229 

230 def to_str(self) -> str: 

231 data = {"arn": self.execution_arn, "seq": self.token_sequence} 

232 json_str = json.dumps(data, separators=(",", ":")) 

233 return base64.b64encode(json_str.encode()).decode() 

234 

235 @classmethod 

236 def from_str(cls, token: str) -> CheckpointToken: 

237 decoded = base64.b64decode(token).decode() 

238 data = json.loads(decoded) 

239 return cls(execution_arn=data["arn"], token_sequence=data["seq"]) 

240 

241 

242@dataclass(frozen=True) 

243class CallbackToken: 

244 """Model a local callback token.""" 

245 

246 execution_arn: str 

247 operation_id: str 

248 

249 def to_str(self) -> str: 

250 data = {"arn": self.execution_arn, "op": self.operation_id} 

251 json_str = json.dumps(data, separators=(",", ":")) 

252 return base64.b64encode(json_str.encode()).decode() 

253 

254 @classmethod 

255 def from_str(cls, token: str) -> CallbackToken: 

256 decoded = base64.b64decode(token).decode() 

257 data = json.loads(decoded) 

258 return cls(execution_arn=data["arn"], operation_id=data["op"]) 

259 

260 

261__all__ = [ 

262 "CallbackToken", 

263 "CheckpointDurableExecutionResponse", 

264 "CheckpointToken", 

265 "GetDurableExecutionStateResponse", 

266 "Invoker", 

267 "LambdaContext", 

268 "SendDurableExecutionCallbackFailureResponse", 

269 "SendDurableExecutionCallbackHeartbeatResponse", 

270 "SendDurableExecutionCallbackSuccessResponse", 

271 "StartDurableExecutionInput", 

272 "StartDurableExecutionOutput", 

273]