Coverage for async_durable_execution/_primitive/base.py: 97%

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

2 

3from __future__ import annotations 

4 

5from abc import ABC, abstractmethod 

6from typing import ClassVar, Generic, TypeVar 

7 

8from .._core import ( 

9 ExecutionState, 

10 InvalidStateError, 

11 Operation, 

12 OperationContext, 

13 OperationIdentifier, 

14 OperationType, 

15 OperationUpdate, 

16 SerDes, 

17 deserialize, 

18 serialize, 

19) 

20 

21T = TypeVar("T") 

22S = TypeVar("S") 

23 

24 

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

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

27 

28 SERDES_OPERATION_TYPE: ClassVar[OperationType | None] = None 

29 

30 def __init__( 

31 self, 

32 state: ExecutionState, 

33 operation_identifier: OperationIdentifier, 

34 ) -> None: 

35 self.state = state 

36 self.operation_identifier = operation_identifier 

37 

38 @property 

39 def operation_id(self) -> str: 

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

41 return self.operation_identifier.require_operation_id() 

42 

43 @property 

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

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

46 return self.operation_identifier.name 

47 

48 @property 

49 def durable_execution_arn(self) -> str: 

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

51 return self.state.durable_execution_arn 

52 

53 async def create_checkpoint( 

54 self, 

55 operation_update: OperationUpdate, 

56 *, 

57 is_sync: bool | None = None, 

58 ) -> Operation | None: 

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

60 if is_sync is None: 

61 return await self.state.create_checkpoint( 

62 operation_update=operation_update, 

63 ) 

64 

65 return await self.state.create_checkpoint( 

66 operation_update=operation_update, 

67 is_sync=is_sync, 

68 ) 

69 

70 async def serialize_value( 

71 self, 

72 value: S, 

73 serdes: SerDes[S] | None, 

74 *, 

75 attempt: int | None = None, 

76 ) -> str: 

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

78 return await serialize( 

79 serdes=serdes, 

80 value=value, 

81 operation_id=self.operation_id, 

82 durable_execution_arn=self.durable_execution_arn, 

83 recursive_level=self.state.recursive_level, 

84 operation_name=self.operation_identifier.name, 

85 parent_id=self.operation_identifier.parent_id, 

86 operation_type=( 

87 self.operation_identifier.operation_type or self.SERDES_OPERATION_TYPE 

88 ), 

89 operation_sub_type=self.operation_identifier.sub_type, 

90 attempt=attempt, 

91 ) 

92 

93 async def deserialize_value( 

94 self, 

95 data: str, 

96 serdes: SerDes[S] | None, 

97 *, 

98 operation: Operation | None = None, 

99 attempt: int | None = None, 

100 ) -> S: 

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

102 return await deserialize( 

103 serdes=serdes, 

104 data=data, 

105 operation_id=self.operation_id, 

106 durable_execution_arn=self.durable_execution_arn, 

107 recursive_level=self.state.recursive_level, 

108 operation_name=( 

109 operation.name 

110 if operation is not None 

111 else self.operation_identifier.name 

112 ), 

113 parent_id=( 

114 operation.parent_id 

115 if operation is not None 

116 else self.operation_identifier.parent_id 

117 ), 

118 operation_type=( 

119 operation.operation_type 

120 if operation is not None 

121 else ( 

122 self.operation_identifier.operation_type 

123 or self.SERDES_OPERATION_TYPE 

124 ) 

125 ), 

126 operation_sub_type=( 

127 operation.sub_type 

128 if operation is not None 

129 else self.operation_identifier.sub_type 

130 ), 

131 attempt=attempt, 

132 ) 

133 

134 @abstractmethod 

135 async def start(self) -> T: 

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

137 ... 

138 

139 @abstractmethod 

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

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

142 ... 

143 

144 async def process(self) -> T: 

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

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

147 if operation is None: 

148 return await self.start() 

149 expected_type = self.operation_identifier.operation_type 

150 if expected_type is not None: 

151 expected = self.operation_identifier 

152 mismatches = [] 

153 if operation.operation_type is not expected_type: 

154 mismatches.append( 

155 f"type={operation.operation_type.value!r}, " 

156 f"expected {expected_type.value!r}" 

157 ) 

158 if operation.sub_type != expected.sub_type: 

159 mismatches.append( 

160 f"sub_type={operation.sub_type!r}, expected {expected.sub_type!r}" 

161 ) 

162 if operation.name != expected.name: 

163 mismatches.append( 

164 f"name={operation.name!r}, expected {expected.name!r}" 

165 ) 

166 if operation.parent_id != expected.parent_id: 

167 mismatches.append( 

168 f"parent_id={operation.parent_id!r}, " 

169 f"expected {expected.parent_id!r}" 

170 ) 

171 if mismatches: 

172 details = "; ".join(mismatches) 

173 msg = ( 

174 f"Reserved extension operation {self.operation_id!r} does " 

175 f"not match its checkpoint: {details}" 

176 ) 

177 raise InvalidStateError(msg) 

178 return await self.replay(operation)