Coverage for async-durable-execution/src/async_durable_execution/primitive/base.py: 100%

1"""Base classes and shared helpers for operation executors.""" 

2 

3from __future__ import annotations 

4 

5from abc import ABC, abstractmethod 

6from dataclasses import dataclass 

7from typing import TYPE_CHECKING, Generic, TypeVar 

8 

9from ..models import ( 

10 Operation, 

11 OperationIdentifier, 

12) 

13from ..serdes import SerDes, deserialize, serialize 

14 

15if TYPE_CHECKING: 

16 from ..models import LambdaContext 

17 from ..models import OperationUpdate 

18 from ..state import ExecutionState 

19 

20T = TypeVar("T") 

21S = TypeVar("S") 

22 

23 

24@dataclass(frozen=True) 

25class OperationContext: 

26 """Protocol defining the interface for durable execution contexts.""" 

27 

28 execution_state: ExecutionState 

29 operation_identifier: OperationIdentifier 

30 

31 @property 

32 def lambda_context(self) -> LambdaContext | None: 

33 """Get the Lambda context for the active invocation.""" 

34 return self.execution_state.lambda_context 

35 

36 @property 

37 def durable_execution_arn(self) -> str: 

38 """Get the ARN of the Durable Execution.""" 

39 return self.execution_state.durable_execution_arn 

40 

41 @property 

42 def parent_id(self) -> str | None: 

43 return self.operation_identifier.parent_id 

44 

45 @property 

46 def operation_id(self) -> str | None: 

47 return self.operation_identifier.operation_id 

48 

49 @property 

50 def operation_name(self) -> str | None: 

51 return self.operation_identifier.name 

52 

53 def is_replaying(self) -> bool: 

54 """Return whether the active context is replaying prior user code.""" 

55 return False 

56 

57 

58class OperationExecutor(ABC, Generic[T]): 

59 """Base class for durable operations with shared state and serdes helpers.""" 

60 

61 def __init__( 

62 self, 

63 state: ExecutionState, 

64 operation_identifier: OperationIdentifier, 

65 ) -> None: 

66 self.state = state 

67 self.operation_identifier = operation_identifier 

68 

69 @property 

70 def operation_id(self) -> str: 

71 """Return the required operation id for this executor.""" 

72 return self.operation_identifier.require_operation_id() 

73 

74 @property 

75 def operation_name(self) -> str | None: 

76 """Return the human-readable operation name, if provided.""" 

77 return self.operation_identifier.name 

78 

79 @property 

80 def durable_execution_arn(self) -> str: 

81 """Return the durable execution ARN for serialization helpers.""" 

82 return self.state.durable_execution_arn 

83 

84 async def create_checkpoint( 

85 self, 

86 operation_update: OperationUpdate, 

87 *, 

88 is_sync: bool | None = None, 

89 ) -> Operation | None: 

90 """Persist a checkpoint update for this operation.""" 

91 if is_sync is None: 

92 return await self.state.create_checkpoint( 

93 operation_update=operation_update, 

94 ) 

95 

96 return await self.state.create_checkpoint( 

97 operation_update=operation_update, 

98 is_sync=is_sync, 

99 ) 

100 

101 async def serialize_value(self, value: S, serdes: SerDes[S] | None) -> str: 

102 """Serialize a value using operation-scoped metadata.""" 

103 return await serialize( 

104 serdes=serdes, 

105 value=value, 

106 operation_id=self.operation_id, 

107 durable_execution_arn=self.durable_execution_arn, 

108 ) 

109 

110 async def deserialize_value(self, data: str, serdes: SerDes[S] | None) -> S: 

111 """Deserialize a value using operation-scoped metadata.""" 

112 return await deserialize( 

113 serdes=serdes, 

114 data=data, 

115 operation_id=self.operation_id, 

116 durable_execution_arn=self.durable_execution_arn, 

117 ) 

118 

119 @abstractmethod 

120 async def start(self) -> T: 

121 """Start a new operation with no existing checkpoint.""" 

122 ... # pragma: no cover 

123 

124 @abstractmethod 

125 async def replay(self, operation: Operation) -> T: 

126 """Replay an operation from an existing checkpoint.""" 

127 ... # pragma: no cover 

128 

129 async def process(self) -> T: 

130 """Process the operation, including replay and checkpoint handling.""" 

131 operation = self.state.operations.get(self.operation_id) 

132 if operation is None: 

133 return await self.start() 

134 return await self.replay(operation)