Coverage for async-durable-execution/src/async_durable_execution/models.py: 99%

359 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-07 16:54 +0000

1from __future__ import annotations 

2 

3import datetime 

4from collections.abc import Mapping, MutableMapping 

5from dataclasses import MISSING, dataclass, field, fields 

6from enum import Enum 

7from typing import ( 

8 Any, 

9 Protocol, 

10 TypeAlias, 

11 get_args, 

12 get_origin, 

13 get_type_hints, 

14) 

15 

16# Replace with `type` it when dropping support to Python 3.11 

17ReplayChildren: TypeAlias = bool 

18OperationPayload: TypeAlias = str 

19TimeoutSeconds: TypeAlias = int 

20 

21 

22class LambdaContext(Protocol): # pragma: no cover 

23 """Minimal AWS Lambda context surface used by the SDK.""" 

24 

25 aws_request_id: str 

26 log_group_name: str | None = None 

27 log_stream_name: str | None = None 

28 function_name: str | None = None 

29 memory_limit_in_mb: str | None = None 

30 function_version: str | None = None 

31 invoked_function_arn: str | None = None 

32 tenant_id: str | None = None 

33 client_context: Any | None = None 

34 identity: Any | None = None 

35 

36 def get_remaining_time_in_millis(self) -> int: ... 

37 def log(self, msg) -> None: ... 

38 

39 

40def _metadata( 

41 *, 

42 alias: str, 

43 serializer: Any = None, 

44 deserializer: Any = None, 

45 omit_if_none: bool = True, 

46 omit_if_falsey: bool = False, 

47 is_timestamp: bool = False, 

48) -> dict[str, Any]: 

49 return { 

50 "alias": alias, 

51 "serializer": serializer, 

52 "deserializer": deserializer, 

53 "omit_if_none": omit_if_none, 

54 "omit_if_falsey": omit_if_falsey, 

55 "is_timestamp": is_timestamp, 

56 } 

57 

58 

59def _enum_type(annotation: Any) -> type[Enum] | None: 

60 if isinstance(annotation, type) and issubclass(annotation, Enum): 

61 return annotation 

62 

63 origin = get_origin(annotation) 

64 if origin is None: 

65 return None 

66 

67 for arg in get_args(annotation): 

68 enum_cls = _enum_type(arg) 

69 if enum_cls is not None: 

70 return enum_cls 

71 

72 return None 

73 

74 

75def _model_type(annotation: Any) -> type[SerializableModel] | None: 

76 if isinstance(annotation, type) and issubclass(annotation, SerializableModel): 

77 return annotation 

78 

79 origin = get_origin(annotation) 

80 if isinstance(origin, type) and issubclass(origin, SerializableModel): 

81 return origin 

82 if origin is None: 

83 return None 

84 

85 for arg in get_args(annotation): 

86 model_cls = _model_type(arg) 

87 if model_cls is not None: 

88 return model_cls 

89 

90 return None 

91 

92 

93def _deserialize_value( 

94 value: Any, 

95 annotation: Any, 

96 metadata: Mapping[str, Any], 

97 *, 

98 json_mode: bool = False, 

99) -> Any: 

100 if value is None: 

101 return None 

102 

103 custom_deserializer = metadata.get("deserializer") 

104 if custom_deserializer is not None: 

105 return custom_deserializer(value) 

106 

107 if json_mode and metadata.get("is_timestamp", False): 

108 return TimestampConverter.from_unix_millis(value) 

109 

110 model_cls = _model_type(annotation) 

111 if model_cls is not None and isinstance(value, Mapping): 

112 if json_mode: 

113 return model_cls.from_json_dict(value) 

114 return model_cls.from_dict(value) 

115 

116 enum_cls = _enum_type(annotation) 

117 if enum_cls is not None: 

118 return enum_cls(value) 

119 

120 origin = get_origin(annotation) 

121 if origin is list: 

122 args = get_args(annotation) 

123 item_annotation = args[0] if args else Any 

124 return [ 

125 _deserialize_value( 

126 item, 

127 item_annotation, 

128 {}, 

129 json_mode=json_mode, 

130 ) 

131 for item in value 

132 ] 

133 

134 return value 

135 

136 

137def _serialize_value( 

138 value: Any, 

139 metadata: Mapping[str, Any], 

140 *, 

141 json_mode: bool = False, 

142) -> Any: 

143 if value is None: 

144 return None 

145 

146 custom_serializer = metadata.get("serializer") 

147 if custom_serializer is not None: 

148 return custom_serializer(value) 

149 

150 if json_mode and metadata.get("is_timestamp", False): 

151 return TimestampConverter.to_unix_millis(value) 

152 

153 if isinstance(value, Enum): 

154 return value.value 

155 

156 if isinstance(value, SerializableModel): 

157 if json_mode: 

158 return value.to_json_dict() 

159 return value.to_dict() 

160 

161 if isinstance(value, list): 

162 return [_serialize_value(item, {}, json_mode=json_mode) for item in value] 

163 

164 return value 

165 

166 

167@dataclass(frozen=True) 

168class SerializableModel: 

169 """Dataclass mixin for the SDK's wire-format serialization helpers.""" 

170 

171 @classmethod 

172 def from_dict(cls, data: Mapping[str, Any]): 

173 return cls._from_mapping(data) 

174 

175 @classmethod 

176 def from_json_dict(cls, data: Mapping[str, Any]): 

177 return cls._from_mapping(data, json_mode=True) 

178 

179 @classmethod 

180 def _from_mapping(cls, data: Mapping[str, Any], *, json_mode: bool = False): 

181 kwargs: dict[str, Any] = {} 

182 type_hints = get_type_hints(cls) 

183 

184 for model_field in fields(cls): 

185 alias = model_field.metadata.get("alias", model_field.name) 

186 annotation = type_hints.get(model_field.name, model_field.type) 

187 

188 if alias in data: 

189 kwargs[model_field.name] = _deserialize_value( 

190 data[alias], 

191 annotation, 

192 model_field.metadata, 

193 json_mode=json_mode, 

194 ) 

195 continue 

196 

197 if ( 

198 model_field.default is not MISSING 

199 or model_field.default_factory is not MISSING 

200 ): 

201 continue 

202 

203 raise KeyError(alias) 

204 

205 return cls(**kwargs) 

206 

207 def to_dict(self) -> MutableMapping[str, Any]: 

208 return self._to_mapping() 

209 

210 def to_json_dict(self) -> MutableMapping[str, Any]: 

211 return self._to_mapping(json_mode=True) 

212 

213 def _to_mapping(self, *, json_mode: bool = False) -> MutableMapping[str, Any]: 

214 result: MutableMapping[str, Any] = {} 

215 

216 for model_field in fields(self): 

217 alias = model_field.metadata.get("alias", model_field.name) 

218 value = getattr(self, model_field.name) 

219 

220 if value is None and model_field.metadata.get("omit_if_none", True): 

221 continue 

222 if not value and model_field.metadata.get("omit_if_falsey", False): 

223 continue 

224 

225 result[alias] = _serialize_value( 

226 value, 

227 model_field.metadata, 

228 json_mode=json_mode, 

229 ) 

230 

231 return result 

232 

233 

234class OperationAction(Enum): 

235 """State transition requested when checkpointing an operation.""" 

236 

237 START = "START" 

238 SUCCEED = "SUCCEED" 

239 FAIL = "FAIL" 

240 RETRY = "RETRY" 

241 CANCEL = "CANCEL" 

242 

243 

244class OperationStatus(Enum): 

245 """Persisted lifecycle status of an operation in execution history.""" 

246 

247 STARTED = "STARTED" 

248 PENDING = "PENDING" 

249 READY = "READY" 

250 SUCCEEDED = "SUCCEEDED" 

251 FAILED = "FAILED" 

252 CANCELLED = "CANCELLED" 

253 TIMED_OUT = "TIMED_OUT" 

254 STOPPED = "STOPPED" 

255 

256 

257class CallbackTimeoutType(Enum): 

258 """Timeout categories surfaced for callback failures.""" 

259 

260 TIMEOUT = "Callback.Timeout" 

261 HEARTBEAT = "Callback.Heartbeat" 

262 

263 

264class ChainedInvokeFailedToStartType(Enum): 

265 """Error type used when a durable invoke never starts remotely.""" 

266 

267 FAILED_TO_START = "ChainedInvoke.FailedToStart" 

268 

269 

270class ChainedInvokeTimeoutType(Enum): 

271 """Error type used when a durable invoke times out.""" 

272 

273 TIMEOUT = "ChainedInvoke.Timeout" 

274 

275 

276class ChainedInvokeStopType(Enum): 

277 """Error type used when a durable invoke is stopped externally.""" 

278 

279 STOPPED = "ChainedInvoke.Stopped" 

280 

281 

282class OperationSubType(Enum): 

283 """Fine-grained operation kind used in execution history.""" 

284 

285 STEP = "Step" 

286 WAIT = "Wait" 

287 CALLBACK = "Callback" 

288 RUN_IN_CHILD_CONTEXT = "RunInChildContext" 

289 MAP = "Map" 

290 MAP_ITERATION = "MapIteration" 

291 PARALLEL = "Parallel" 

292 PARALLEL_BRANCH = "ParallelBranch" 

293 WAIT_FOR_CALLBACK = "WaitForCallback" 

294 WAIT_FOR_CONDITION = "WaitForCondition" 

295 CHAINED_INVOKE = "ChainedInvoke" 

296 EXECUTION = "Execution" 

297 

298 

299class OperationType(Enum): 

300 """Top-level operation categories persisted by the durable backend.""" 

301 

302 EXECUTION = "EXECUTION" 

303 CONTEXT = "CONTEXT" 

304 STEP = "STEP" 

305 WAIT = "WAIT" 

306 CALLBACK = "CALLBACK" 

307 CHAINED_INVOKE = "CHAINED_INVOKE" 

308 

309 @classmethod 

310 def from_sub_type(cls, sub_type: OperationSubType) -> OperationType: 

311 match sub_type: 

312 case OperationSubType.STEP | OperationSubType.WAIT_FOR_CONDITION: 

313 return OperationType.STEP 

314 case OperationSubType.WAIT: 

315 return OperationType.WAIT 

316 case OperationSubType.CHAINED_INVOKE: 

317 return OperationType.CHAINED_INVOKE 

318 case OperationSubType.CALLBACK: 

319 return OperationType.CALLBACK 

320 case OperationSubType.EXECUTION: 

321 return OperationType.EXECUTION 

322 case ( 

323 OperationSubType.WAIT_FOR_CALLBACK 

324 | OperationSubType.RUN_IN_CHILD_CONTEXT 

325 | OperationSubType.MAP 

326 | OperationSubType.MAP_ITERATION 

327 | OperationSubType.PARALLEL 

328 | OperationSubType.PARALLEL_BRANCH 

329 ): 

330 return OperationType.CONTEXT 

331 case _: 

332 raise ValueError(f"Unknown operation sub-type {sub_type}") 

333 

334 

335@dataclass(frozen=True) 

336class OperationIdentifier: 

337 """Container for operation id, parent id, and name.""" 

338 

339 operation_id: str | None 

340 sub_type: OperationSubType 

341 parent_id: str | None = None 

342 name: str | None = None 

343 

344 @property 

345 def type(self) -> OperationType: 

346 return OperationType.from_sub_type(self.sub_type) 

347 

348 def require_operation_id(self) -> str: 

349 """Return the operation id for non-root operations.""" 

350 if self.operation_id is None: 

351 msg = "operation_id is required for non-execution operations" 

352 raise ValueError(msg) 

353 return self.operation_id 

354 

355 @classmethod 

356 def create_execution_op(cls): 

357 return cls(None, OperationSubType.EXECUTION, None, None) 

358 

359 

360class InvocationStatus(Enum): 

361 """Overall result of a single durable Lambda invocation.""" 

362 

363 SUCCEEDED = "SUCCEEDED" 

364 FAILED = "FAILED" 

365 PENDING = "PENDING" 

366 

367 # Used internally only: the invocation failed and the backend will retry 

368 RETRY = "RETRY" 

369 

370 

371@dataclass(frozen=True) 

372class ErrorObject(SerializableModel): 

373 """Serializable representation of an exception captured by the SDK.""" 

374 

375 message: str | None = field(default=None, metadata=_metadata(alias="ErrorMessage")) 

376 type: str | None = field(default=None, metadata=_metadata(alias="ErrorType")) 

377 data: str | None = field(default=None, metadata=_metadata(alias="ErrorData")) 

378 stack_trace: list[str] | None = field( 

379 default=None, metadata=_metadata(alias="StackTrace") 

380 ) 

381 

382 @classmethod 

383 def from_exception(cls, exception: Exception) -> ErrorObject: 

384 return cls( 

385 message=str(exception), 

386 type=type(exception).__name__, 

387 data=None, 

388 stack_trace=None, 

389 ) 

390 

391 @classmethod 

392 def from_message(cls, message: str) -> ErrorObject: 

393 return cls( 

394 message=message, 

395 type=None, 

396 data=None, 

397 stack_trace=None, 

398 ) 

399 

400 

401@dataclass(frozen=True) 

402class DurableExecutionInvocationOutput(SerializableModel): 

403 """Representation the DurableExecutionInvocationOutput. This is what the Durable lambda handler returns. 

404 

405 If the execution has been already completed via an update to the EXECUTION operation via CheckpointDurableExecution, 

406 payload must be empty for SUCCEEDED/FAILED status. 

407 """ 

408 

409 status: InvocationStatus = field(metadata=_metadata(alias="Status")) 

410 result: str | None = field(default=None, metadata=_metadata(alias="Result")) 

411 error: ErrorObject | None = field(default=None, metadata=_metadata(alias="Error")) 

412 

413 @classmethod 

414 def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput: 

415 return cls(status=InvocationStatus.SUCCEEDED, result=result) 

416 

417 @classmethod 

418 def create_retry(cls, error: ErrorObject) -> DurableExecutionInvocationOutput: 

419 return cls(status=InvocationStatus.RETRY, error=error) 

420 

421 

422@dataclass(frozen=True) 

423class ExecutionDetails(SerializableModel): 

424 """Extra fields stored on the root execution operation.""" 

425 

426 input_payload: str | None = field( 

427 default=None, 

428 metadata=_metadata(alias="InputPayload", omit_if_none=False), 

429 ) 

430 

431 

432@dataclass(frozen=True) 

433class ContextDetails(SerializableModel): 

434 """Checkpoint payload stored for child-context style operations.""" 

435 

436 replay_children: ReplayChildren = field( 

437 default=False, metadata=_metadata(alias="ReplayChildren") 

438 ) 

439 result: OperationPayload | None = field( 

440 default=None, metadata=_metadata(alias="Result") 

441 ) 

442 error: ErrorObject | None = field(default=None, metadata=_metadata(alias="Error")) 

443 

444 

445@dataclass(frozen=True) 

446class StepDetails(SerializableModel): 

447 """Checkpoint payload stored for durable steps and polling checks.""" 

448 

449 attempt: int = field(default=0, metadata=_metadata(alias="Attempt")) 

450 next_attempt_timestamp: datetime.datetime | None = field( 

451 default=None, 

452 metadata=_metadata(alias="NextAttemptTimestamp", is_timestamp=True), 

453 ) 

454 result: OperationPayload | None = field( 

455 default=None, 

456 metadata=_metadata(alias="Result"), 

457 ) 

458 error: ErrorObject | None = field( 

459 default=None, 

460 metadata=_metadata(alias="Error"), 

461 ) 

462 

463 

464@dataclass(frozen=True) 

465class WaitDetails(SerializableModel): 

466 """Checkpoint payload stored for durable waits.""" 

467 

468 scheduled_end_timestamp: datetime.datetime | None = field( 

469 default=None, 

470 metadata=_metadata(alias="ScheduledEndTimestamp", is_timestamp=True), 

471 ) 

472 

473 

474@dataclass(frozen=True) 

475class CallbackDetails(SerializableModel): 

476 """Checkpoint payload stored for callbacks and callback results.""" 

477 

478 callback_id: str = field(metadata=_metadata(alias="CallbackId")) 

479 result: str | None = field(default=None, metadata=_metadata(alias="Result")) 

480 error: ErrorObject | None = field(default=None, metadata=_metadata(alias="Error")) 

481 

482 

483@dataclass(frozen=True) 

484class ChainedInvokeDetails(SerializableModel): 

485 """Checkpoint payload stored for durable invokes.""" 

486 

487 result: str | None = field(default=None, metadata=_metadata(alias="Result")) 

488 error: ErrorObject | None = field(default=None, metadata=_metadata(alias="Error")) 

489 

490 

491@dataclass(frozen=True) 

492class StepOptions(SerializableModel): 

493 """Additional options recorded on step retries.""" 

494 

495 next_attempt_delay_seconds: int = field( 

496 default=0, 

497 metadata=_metadata(alias="NextAttemptDelaySeconds"), 

498 ) 

499 

500 

501@dataclass(frozen=True) 

502class WaitOptions(SerializableModel): 

503 """ 

504 Wait Options provides details regarding suspension. 

505 

506 As of 2025/10/27: 

507 

508 - `wait_seconds` accepts values between 1, and 31622400 

509 - When wait_second seconds does not exist,then we default to 1 

510 

511 """ 

512 

513 wait_seconds: int = field(default=1, metadata=_metadata(alias="WaitSeconds")) 

514 

515 

516@dataclass(frozen=True) 

517class CallbackOptions(SerializableModel): 

518 """ 

519 Callback options provides details about the callback, wrt timeout 

520 and heartbeat checks. 

521 

522 As of 2025/10/27: 

523 - When timeout_seconds == 0, then the callback has no timeout 

524 - When heartbeat_timeout_seconds == 0, then the callback has no timeout 

525 

526 - When timeout_seconds is not present, then default is 0 

527 - When heartbeat_timeout_seconds, then default is 0 

528 

529 """ 

530 

531 timeout_seconds: TimeoutSeconds = field( 

532 default=0, metadata=_metadata(alias="TimeoutSeconds") 

533 ) 

534 heartbeat_timeout_seconds: int = field( 

535 default=0, 

536 metadata=_metadata(alias="HeartbeatTimeoutSeconds"), 

537 ) 

538 

539 

540@dataclass(frozen=True) 

541class ChainedInvokeOptions(SerializableModel): 

542 """ 

543 As of 2025/10/27: 

544 - Chained invoke options only contains a function name 

545 """ 

546 

547 function_name: str = field(metadata=_metadata(alias="FunctionName")) 

548 tenant_id: str | None = field(default=None, metadata=_metadata(alias="TenantId")) 

549 

550 

551@dataclass(frozen=True) 

552class ContextOptions(SerializableModel): 

553 """Extra flags recorded for child-context operations.""" 

554 

555 replay_children: ReplayChildren = field( 

556 default=False, 

557 metadata=_metadata(alias="ReplayChildren"), 

558 ) 

559 

560 

561@dataclass(frozen=True) 

562class OperationUpdate(SerializableModel): 

563 """Update an Operation. Use this to create a checkpoint. 

564 

565 See the various create_ factory class methods to instantiate me. 

566 """ 

567 

568 operation_id: str = field(metadata=_metadata(alias="Id")) 

569 operation_type: OperationType = field(metadata=_metadata(alias="Type")) 

570 action: OperationAction = field(metadata=_metadata(alias="Action")) 

571 parent_id: str | None = field( 

572 default=None, 

573 metadata=_metadata(alias="ParentId", omit_if_falsey=True), 

574 ) 

575 name: str | None = field( 

576 default=None, 

577 metadata=_metadata(alias="Name", omit_if_falsey=True), 

578 ) 

579 sub_type: OperationSubType | None = field( 

580 default=None, 

581 metadata=_metadata(alias="SubType"), 

582 ) 

583 payload: str | None = field( 

584 default=None, 

585 metadata=_metadata(alias="Payload", omit_if_falsey=True), 

586 ) 

587 error: ErrorObject | None = field(default=None, metadata=_metadata(alias="Error")) 

588 context_options: ContextOptions | None = field( 

589 default=None, 

590 metadata=_metadata(alias="ContextOptions"), 

591 ) 

592 step_options: StepOptions | None = field( 

593 default=None, 

594 metadata=_metadata(alias="StepOptions"), 

595 ) 

596 wait_options: WaitOptions | None = field( 

597 default=None, 

598 metadata=_metadata(alias="WaitOptions"), 

599 ) 

600 callback_options: CallbackOptions | None = field( 

601 default=None, 

602 metadata=_metadata(alias="CallbackOptions"), 

603 ) 

604 chained_invoke_options: ChainedInvokeOptions | None = field( 

605 default=None, 

606 metadata=_metadata(alias="ChainedInvokeOptions"), 

607 ) 

608 

609 @classmethod 

610 def create_callback( 

611 cls, identifier: OperationIdentifier, callback_options: CallbackOptions 

612 ) -> OperationUpdate: 

613 """Create an instance of OperationUpdate for type:CALLBACK, action:START""" 

614 return cls( 

615 operation_id=identifier.require_operation_id(), 

616 parent_id=identifier.parent_id, 

617 operation_type=OperationType.CALLBACK, 

618 sub_type=OperationSubType.CALLBACK, 

619 action=OperationAction.START, 

620 name=identifier.name, 

621 callback_options=callback_options, 

622 ) 

623 

624 @classmethod 

625 def create_context_start( 

626 cls, identifier: OperationIdentifier, sub_type: OperationSubType 

627 ) -> OperationUpdate: 

628 """Create an instance of OperationUpdate for type: CONTEXT, action: START.""" 

629 return cls( 

630 operation_id=identifier.require_operation_id(), 

631 parent_id=identifier.parent_id, 

632 operation_type=OperationType.CONTEXT, 

633 sub_type=sub_type, 

634 action=OperationAction.START, 

635 name=identifier.name, 

636 ) 

637 

638 @classmethod 

639 def create_context_succeed( 

640 cls, 

641 identifier: OperationIdentifier, 

642 payload: str, 

643 sub_type: OperationSubType, 

644 context_options: ContextOptions | None = None, 

645 ) -> OperationUpdate: 

646 """Create an instance of OperationUpdate for type: CONTEXT, action: SUCCEED.""" 

647 return cls( 

648 operation_id=identifier.require_operation_id(), 

649 parent_id=identifier.parent_id, 

650 operation_type=OperationType.CONTEXT, 

651 sub_type=sub_type, 

652 action=OperationAction.SUCCEED, 

653 name=identifier.name, 

654 payload=payload, 

655 context_options=context_options, 

656 ) 

657 

658 @classmethod 

659 def create_context_fail( 

660 cls, 

661 identifier: OperationIdentifier, 

662 error: ErrorObject, 

663 sub_type: OperationSubType, 

664 ) -> OperationUpdate: 

665 """Create an instance of OperationUpdate for type: CONTEXT, action: FAIL.""" 

666 return cls( 

667 operation_id=identifier.require_operation_id(), 

668 parent_id=identifier.parent_id, 

669 operation_type=OperationType.CONTEXT, 

670 sub_type=sub_type, 

671 action=OperationAction.FAIL, 

672 name=identifier.name, 

673 error=error, 

674 ) 

675 

676 @classmethod 

677 def create_execution_succeed(cls, payload: str) -> OperationUpdate: 

678 """Create an instance of OperationUpdate for type: EXECUTION, action: SUCCEED.""" 

679 return cls( 

680 operation_id=f"execution-result-{int(datetime.datetime.now(tz=datetime.timezone.utc).timestamp() * 1000)}", 

681 operation_type=OperationType.EXECUTION, 

682 action=OperationAction.SUCCEED, 

683 payload=payload, 

684 ) 

685 

686 @classmethod 

687 def create_execution_fail(cls, error: ErrorObject) -> OperationUpdate: 

688 """Create an instance of OperationUpdate for type: EXECUTION, action: FAIL.""" 

689 return cls( 

690 operation_id=f"execution-result-{int(datetime.datetime.now(tz=datetime.timezone.utc).timestamp() * 1000)}", 

691 operation_type=OperationType.EXECUTION, 

692 action=OperationAction.FAIL, 

693 error=error, 

694 ) 

695 

696 @classmethod 

697 def create_step_succeed( 

698 cls, identifier: OperationIdentifier, payload: str 

699 ) -> OperationUpdate: 

700 """Create an instance of OperationUpdate for type: STEP, action: SUCCEED.""" 

701 return cls( 

702 operation_id=identifier.require_operation_id(), 

703 parent_id=identifier.parent_id, 

704 operation_type=OperationType.STEP, 

705 sub_type=OperationSubType.STEP, 

706 action=OperationAction.SUCCEED, 

707 name=identifier.name, 

708 payload=payload, 

709 ) 

710 

711 @classmethod 

712 def create_step_fail( 

713 cls, identifier: OperationIdentifier, error: ErrorObject 

714 ) -> OperationUpdate: 

715 """Create an instance of OperationUpdate for type: STEP, action: FAIL.""" 

716 return cls( 

717 operation_id=identifier.require_operation_id(), 

718 parent_id=identifier.parent_id, 

719 operation_type=OperationType.STEP, 

720 sub_type=OperationSubType.STEP, 

721 action=OperationAction.FAIL, 

722 name=identifier.name, 

723 error=error, 

724 ) 

725 

726 @classmethod 

727 def create_step_start(cls, identifier: OperationIdentifier) -> OperationUpdate: 

728 """Create an instance of OperationUpdate for type: STEP, action: START.""" 

729 return cls( 

730 operation_id=identifier.require_operation_id(), 

731 parent_id=identifier.parent_id, 

732 operation_type=OperationType.STEP, 

733 sub_type=OperationSubType.STEP, 

734 action=OperationAction.START, 

735 name=identifier.name, 

736 ) 

737 

738 @classmethod 

739 def create_step_retry( 

740 cls, 

741 identifier: OperationIdentifier, 

742 error: ErrorObject, 

743 next_attempt_delay_seconds: int, 

744 ) -> OperationUpdate: 

745 """Create an instance of OperationUpdate for type: STEP, action: RETRY.""" 

746 return cls( 

747 operation_id=identifier.require_operation_id(), 

748 parent_id=identifier.parent_id, 

749 operation_type=OperationType.STEP, 

750 sub_type=OperationSubType.STEP, 

751 action=OperationAction.RETRY, 

752 name=identifier.name, 

753 error=error, 

754 step_options=StepOptions( 

755 next_attempt_delay_seconds=next_attempt_delay_seconds 

756 ), 

757 ) 

758 

759 @classmethod 

760 def create_invoke_start( 

761 cls, 

762 identifier: OperationIdentifier, 

763 payload: str, 

764 chained_invoke_options: ChainedInvokeOptions, 

765 ) -> OperationUpdate: 

766 """Create an instance of OperationUpdate for type: INVOKE, action: START.""" 

767 return cls( 

768 operation_id=identifier.require_operation_id(), 

769 parent_id=identifier.parent_id, 

770 operation_type=OperationType.CHAINED_INVOKE, 

771 sub_type=OperationSubType.CHAINED_INVOKE, 

772 action=OperationAction.START, 

773 name=identifier.name, 

774 payload=payload, 

775 chained_invoke_options=chained_invoke_options, 

776 ) 

777 

778 @classmethod 

779 def create_wait_for_condition_start( 

780 cls, identifier: OperationIdentifier 

781 ) -> OperationUpdate: 

782 """Create an instance of OperationUpdate for type: STEP, action: START.""" 

783 return cls( 

784 operation_id=identifier.require_operation_id(), 

785 parent_id=identifier.parent_id, 

786 operation_type=OperationType.STEP, 

787 sub_type=OperationSubType.WAIT_FOR_CONDITION, 

788 action=OperationAction.START, 

789 name=identifier.name, 

790 ) 

791 

792 @classmethod 

793 def create_wait_for_condition_succeed( 

794 cls, identifier: OperationIdentifier, payload: str 

795 ) -> OperationUpdate: 

796 """Create an instance of OperationUpdate for type: STEP, action: SUCCEED.""" 

797 return cls( 

798 operation_id=identifier.require_operation_id(), 

799 parent_id=identifier.parent_id, 

800 operation_type=OperationType.STEP, 

801 sub_type=OperationSubType.WAIT_FOR_CONDITION, 

802 action=OperationAction.SUCCEED, 

803 name=identifier.name, 

804 payload=payload, 

805 ) 

806 

807 @classmethod 

808 def create_wait_for_condition_retry( 

809 cls, 

810 identifier: OperationIdentifier, 

811 payload: str, 

812 next_attempt_delay_seconds: int, 

813 ) -> OperationUpdate: 

814 """Create an instance of OperationUpdate for type: STEP, action: RETRY.""" 

815 return cls( 

816 operation_id=identifier.require_operation_id(), 

817 parent_id=identifier.parent_id, 

818 operation_type=OperationType.STEP, 

819 sub_type=OperationSubType.WAIT_FOR_CONDITION, 

820 action=OperationAction.RETRY, 

821 name=identifier.name, 

822 payload=payload, 

823 step_options=StepOptions( 

824 next_attempt_delay_seconds=next_attempt_delay_seconds 

825 ), 

826 ) 

827 

828 @classmethod 

829 def create_wait_for_condition_fail( 

830 cls, identifier: OperationIdentifier, error: ErrorObject 

831 ) -> OperationUpdate: 

832 """Create an instance of OperationUpdate for type: STEP, action: FAIL.""" 

833 return cls( 

834 operation_id=identifier.require_operation_id(), 

835 parent_id=identifier.parent_id, 

836 operation_type=OperationType.STEP, 

837 sub_type=OperationSubType.WAIT_FOR_CONDITION, 

838 action=OperationAction.FAIL, 

839 name=identifier.name, 

840 error=error, 

841 ) 

842 

843 @classmethod 

844 def create_wait_start( 

845 cls, identifier: OperationIdentifier, wait_options: WaitOptions 

846 ) -> OperationUpdate: 

847 """Create an instance of OperationUpdate for type: WAIT, action: START.""" 

848 return cls( 

849 operation_id=identifier.require_operation_id(), 

850 parent_id=identifier.parent_id, 

851 operation_type=OperationType.WAIT, 

852 sub_type=OperationSubType.WAIT, 

853 action=OperationAction.START, 

854 name=identifier.name, 

855 wait_options=wait_options, 

856 ) 

857 

858 

859class TimestampConverter: 

860 """Converter for datetime/Unix timestamp conversions.""" 

861 

862 @staticmethod 

863 def to_unix_millis(dt: datetime.datetime | None) -> int | None: 

864 """Convert datetime to Unix timestamp in milliseconds.""" 

865 return int(dt.timestamp() * 1000) if dt else None 

866 

867 @staticmethod 

868 def from_unix_millis(ms: int | None) -> datetime.datetime | None: 

869 """Convert Unix timestamp in milliseconds to datetime.""" 

870 return ( 

871 datetime.datetime.fromtimestamp(ms / 1000, tz=datetime.timezone.utc) 

872 if ms is not None 

873 else None 

874 ) 

875 

876 

877@dataclass(frozen=True) 

878class Operation(SerializableModel): 

879 """Represent the Operation type for GetDurableExecutionState and CheckpointDurableExecution.""" 

880 

881 operation_id: str = field(metadata=_metadata(alias="Id")) 

882 operation_type: OperationType = field(metadata=_metadata(alias="Type")) 

883 status: OperationStatus = field(metadata=_metadata(alias="Status")) 

884 parent_id: str | None = field( 

885 default=None, 

886 metadata=_metadata(alias="ParentId", omit_if_falsey=True), 

887 ) 

888 name: str | None = field( 

889 default=None, 

890 metadata=_metadata(alias="Name", omit_if_falsey=True), 

891 ) 

892 start_timestamp: datetime.datetime | None = field( 

893 default=None, 

894 metadata=_metadata(alias="StartTimestamp", is_timestamp=True), 

895 ) 

896 end_timestamp: datetime.datetime | None = field( 

897 default=None, 

898 metadata=_metadata(alias="EndTimestamp", is_timestamp=True), 

899 ) 

900 sub_type: OperationSubType | None = field( 

901 default=None, 

902 metadata=_metadata(alias="SubType"), 

903 ) 

904 execution_details: ExecutionDetails | None = field( 

905 default=None, 

906 metadata=_metadata(alias="ExecutionDetails"), 

907 ) 

908 context_details: ContextDetails | None = field( 

909 default=None, 

910 metadata=_metadata(alias="ContextDetails"), 

911 ) 

912 step_details: StepDetails | None = field( 

913 default=None, 

914 metadata=_metadata(alias="StepDetails"), 

915 ) 

916 wait_details: WaitDetails | None = field( 

917 default=None, 

918 metadata=_metadata(alias="WaitDetails"), 

919 ) 

920 callback_details: CallbackDetails | None = field( 

921 default=None, 

922 metadata=_metadata(alias="CallbackDetails"), 

923 ) 

924 chained_invoke_details: ChainedInvokeDetails | None = field( 

925 default=None, 

926 metadata=_metadata(alias="ChainedInvokeDetails"), 

927 ) 

928 

929 

930@dataclass(frozen=True) 

931class CheckpointUpdatedExecutionState(SerializableModel): 

932 """Representation of the CheckpointUpdatedExecutionState structure of the DEX API.""" 

933 

934 operations: list[Operation] = field( 

935 default_factory=list, 

936 metadata=_metadata(alias="Operations"), 

937 ) 

938 next_marker: str | None = field( 

939 default=None, metadata=_metadata(alias="NextMarker") 

940 ) 

941 

942 

943@dataclass(frozen=True) 

944class CheckpointOutput(SerializableModel): 

945 """Representation of the CheckpointDurableExecutionOutput structure of the DEX CheckpointDurableExecution API.""" 

946 

947 checkpoint_token: str = field( 

948 default="", 

949 metadata=_metadata(alias="CheckpointToken", omit_if_none=False), 

950 ) 

951 new_execution_state: CheckpointUpdatedExecutionState = field( 

952 default_factory=CheckpointUpdatedExecutionState, 

953 metadata=_metadata(alias="NewExecutionState", omit_if_none=False), 

954 ) 

955 

956 

957@dataclass(frozen=True) 

958class StateOutput(SerializableModel): 

959 """Representation of the GetDurableExecutionStateOutput structure of the DEX GetDurableExecutionState API.""" 

960 

961 operations: list[Operation] = field( 

962 default_factory=list, 

963 metadata=_metadata(alias="Operations"), 

964 ) 

965 next_marker: str | None = field( 

966 default=None, metadata=_metadata(alias="NextMarker") 

967 ) 

968 

969 

970__all__ = [ 

971 "CallbackDetails", 

972 "CallbackOptions", 

973 "CallbackTimeoutType", 

974 "ChainedInvokeDetails", 

975 "ChainedInvokeFailedToStartType", 

976 "ChainedInvokeOptions", 

977 "ChainedInvokeStopType", 

978 "ChainedInvokeTimeoutType", 

979 "CheckpointOutput", 

980 "CheckpointUpdatedExecutionState", 

981 "ContextDetails", 

982 "ContextOptions", 

983 "DurableExecutionInvocationOutput", 

984 "ErrorObject", 

985 "ExecutionDetails", 

986 "InvocationStatus", 

987 "Operation", 

988 "OperationAction", 

989 "OperationPayload", 

990 "OperationStatus", 

991 "OperationSubType", 

992 "OperationType", 

993 "OperationUpdate", 

994 "ReplayChildren", 

995 "SerializableModel", 

996 "StateOutput", 

997 "StepDetails", 

998 "StepOptions", 

999 "TimeoutSeconds", 

1000 "TimestampConverter", 

1001 "WaitDetails", 

1002 "WaitOptions", 

1003 "OperationIdentifier", 

1004]