Coverage for async_durable_execution/_primitive/callback.py: 95%
81 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Implementation for the backend-supported create_callback operation."""
3from __future__ import annotations
5import asyncio
6import logging
7from typing import Any, Generic, TypeVar
9from .base import OperationExecutor
10from .._core import (
11 CallbackOptions,
12 CallbackTimeoutType,
13 Duration,
14 DurableContext,
15 ExecutionError,
16 ExecutionState,
17 Operation,
18 OperationIdentifier,
19 OperationStatus,
20 OperationType,
21 OperationUpdate,
22 PassThroughSerDes,
23 SerDes,
24 SuspendExecution,
25 TerminationReason,
26 _register_sdk_control_error_type,
27 deserialize,
28 duration_to_seconds,
29)
31T = TypeVar("T") # Result type
33logger = logging.getLogger(__name__)
35PASS_THROUGH_SERDES: SerDes[Any] = PassThroughSerDes()
36_LEGACY_CALLBACK_ERROR_TYPE_NAMES = (
37 "async_durable_execution.exceptions.CallbackError",
38 "async_durable_execution.primitive.callback.CallbackError",
39)
42class CallbackError(ExecutionError):
43 """Error in callback handling."""
45 def __init__(self, message: str, callback_id: str | None = None) -> None:
46 super().__init__(message, TerminationReason.CALLBACK_ERROR)
47 self.callback_id = callback_id
50def _encode_callback_error_payload(error: ExecutionError) -> str | None:
51 if not isinstance(error, CallbackError): 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true
52 msg = "CallbackError codec received an incompatible exception."
53 raise TypeError(msg)
54 return error.callback_id
57def _restore_callback_error(
58 message: str,
59 payload: str | None,
60) -> CallbackError:
61 return CallbackError(message=message, callback_id=payload)
64_register_sdk_control_error_type(
65 CallbackError,
66 encode_payload=_encode_callback_error_payload,
67 restore=_restore_callback_error,
68 legacy_exception_type_names=_LEGACY_CALLBACK_ERROR_TYPE_NAMES,
69)
72class CallbackOperationExecutor(OperationExecutor[str]):
73 """Executor for callback operations."""
75 SERDES_OPERATION_TYPE = OperationType.CALLBACK
77 def __init__(
78 self,
79 state: ExecutionState,
80 operation_identifier: OperationIdentifier,
81 timeout: Duration | None = None,
82 heartbeat_timeout: Duration | None = None,
83 ) -> None:
84 """Initialize the callback operation executor.
86 Args:
87 state: The execution state
88 operation_identifier: The operation identifier
89 timeout: Optional maximum time to wait for callback completion.
90 heartbeat_timeout: Optional maximum time to wait between callback heartbeats.
91 """
92 super().__init__(state=state, operation_identifier=operation_identifier)
93 self.timeout_seconds = (
94 duration_to_seconds(timeout, "timeout") if timeout is not None else 0
95 )
96 self.heartbeat_timeout_seconds = (
97 duration_to_seconds(heartbeat_timeout, "heartbeat_timeout")
98 if heartbeat_timeout is not None
99 else 0
100 )
102 async def start(self) -> str:
103 """Start a new callback operation."""
104 callback_options = CallbackOptions(
105 timeout_seconds=self.timeout_seconds,
106 heartbeat_timeout_seconds=self.heartbeat_timeout_seconds,
107 )
109 create_callback_operation: OperationUpdate = OperationUpdate.create_callback(
110 identifier=self.operation_identifier,
111 callback_options=callback_options,
112 )
114 operation = await self.create_checkpoint(create_callback_operation)
116 if not operation:
117 msg = f"Missing callback details for operation: {self.operation_identifier.operation_id}"
118 raise CallbackError(msg)
119 return await self.replay(operation)
121 async def replay(self, operation: Operation) -> str:
122 """Replay an existing callback operation from its checkpoint."""
123 if not operation.callback_details:
124 msg = (
125 f"Missing callback details for operation: "
126 f"{self.operation_identifier.operation_id}"
127 )
128 raise CallbackError(msg)
130 return await self.execute(operation)
132 async def execute(self, operation: Operation | None) -> str:
133 """Execute callback operation by extracting the callback_id.
135 Callbacks don't execute logic - they just extract and return the callback_id
136 from the operation data.
138 Args:
139 operation: The callback operation containing callback_details
141 Returns:
142 The callback_id from the checkpoint
144 Raises:
145 CallbackError: If callback_details are missing (should never happen)
146 """
147 if operation is None or not operation.callback_details:
148 msg = f"Missing callback details for operation: {self.operation_identifier.operation_id}"
149 raise CallbackError(msg)
151 return operation.callback_details.callback_id
154def create_callback(
155 *,
156 name: str | None = None,
157 timeout: Duration | None = None,
158 heartbeat_timeout: Duration | None = None,
159 serdes: SerDes | None = None,
160) -> asyncio.Task[Callback]:
161 """Compatibility import for the canonical operation-layer helper."""
162 from .._operation.callback import create_callback as operation_create_callback
164 return operation_create_callback(
165 name=name,
166 timeout=timeout,
167 heartbeat_timeout=heartbeat_timeout,
168 serdes=serdes,
169 )
172async def _create_callback(
173 *,
174 context: DurableContext,
175 operation_identifier: OperationIdentifier,
176 operation_id: str,
177 timeout: Duration | None = None,
178 heartbeat_timeout: Duration | None = None,
179 serdes: SerDes | None = None,
180) -> Callback:
181 executor: CallbackOperationExecutor = CallbackOperationExecutor(
182 state=context.execution_state,
183 operation_identifier=operation_identifier,
184 timeout=timeout,
185 heartbeat_timeout=heartbeat_timeout,
186 )
187 callback_id: str = await executor.process()
188 return Callback(
189 callback_id=callback_id,
190 operation_id=operation_id,
191 state=context.execution_state,
192 serdes=serdes,
193 )
196class Callback(Generic[T]):
197 """A future that will block on result() until callback_id returns."""
199 def __init__(
200 self,
201 callback_id: str,
202 operation_id: str,
203 state: ExecutionState,
204 serdes: SerDes[T] | None = None,
205 ) -> None:
206 self.callback_id: str = callback_id
207 self.operation_id: str = operation_id
208 self.state: ExecutionState = state
209 self.serdes: SerDes[T] | None = serdes
211 async def result(self) -> T | None:
212 """Return the result of the future. Will block until result is available.
214 This will suspend the current execution while waiting for the result to
215 become available. Durable Functions will replay the execution once the
216 result is ready, and proceed when it reaches the .result() call.
218 Use the callback id with the following APIs to send back the result, error or
219 heartbeats: SendDurableExecutionCallbackSuccess, SendDurableExecutionCallbackFailure
220 and SendDurableExecutionCallbackHeartbeat.
221 """
222 operation = self.state.operations.get(self.operation_id)
224 if not isinstance(operation, Operation):
225 msg = "Callback operation must exist"
226 raise CallbackError(message=msg, callback_id=self.callback_id)
228 if operation.status in {
229 OperationStatus.FAILED,
230 OperationStatus.CANCELLED,
231 OperationStatus.TIMED_OUT,
232 OperationStatus.STOPPED,
233 }:
234 msg = _format_callback_error_message(operation)
235 raise CallbackError(message=msg, callback_id=self.callback_id)
237 if operation.status is OperationStatus.SUCCEEDED:
238 if (
239 not operation.callback_details
240 or operation.callback_details.result is None
241 ):
242 return None
244 return await deserialize(
245 serdes=self.serdes if self.serdes is not None else PASS_THROUGH_SERDES,
246 data=operation.callback_details.result,
247 operation_id=self.operation_id,
248 durable_execution_arn=self.state.durable_execution_arn,
249 recursive_level=self.state.recursive_level,
250 operation_name=operation.name,
251 parent_id=operation.parent_id,
252 operation_type=operation.operation_type,
253 operation_sub_type=operation.sub_type,
254 )
256 # operation exists; it has not terminated (successfully or otherwise)
257 # therefore we should wait
258 msg = "Callback result not received yet. Suspending execution while waiting for result."
259 raise SuspendExecution(msg)
262def _format_callback_error_message(operation: Operation) -> str:
263 """Build a stable callback error message from checkpoint state."""
264 error = operation.callback_details.error if operation.callback_details else None
265 if not error or not error.message:
266 return "Callback failed"
268 message = error.message
269 if (
270 operation.status is OperationStatus.TIMED_OUT
271 and error.type in {timeout.value for timeout in CallbackTimeoutType}
272 and error.type not in message
273 ):
274 return f"{message}: {error.type}"
276 return message