Coverage for async_durable_execution/_primitive/wait.py: 94%

1"""Implement the durable wait operation.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7 

8from .base import OperationExecutor 

9from .._core import ( 

10 Duration, 

11 DurableContext, 

12 ExecutionState, 

13 Operation, 

14 OperationIdentifier, 

15 OperationStatus, 

16 OperationType, 

17 OperationUpdate, 

18 WaitOptions, 

19 suspend_with_optional_resume_delay, 

20) 

21 

22logger = logging.getLogger(__name__) 

23 

24 

25class WaitOperationExecutor(OperationExecutor[None]): 

26 """Executor for wait operations.""" 

27 

28 SERDES_OPERATION_TYPE = OperationType.WAIT 

29 

30 def __init__( 

31 self, 

32 seconds: int, 

33 state: ExecutionState, 

34 operation_identifier: OperationIdentifier, 

35 ) -> None: 

36 """Initialize the wait operation executor. 

37 

38 Args: 

39 seconds: Number of seconds to wait 

40 state: The execution state 

41 operation_identifier: The operation identifier 

42 """ 

43 super().__init__(state=state, operation_identifier=operation_identifier) 

44 self.seconds = seconds 

45 

46 async def start(self) -> None: 

47 """Start a new wait operation.""" 

48 operation: OperationUpdate = OperationUpdate.create_wait_start( 

49 identifier=self.operation_identifier, 

50 wait_options=WaitOptions(wait_seconds=self.seconds), 

51 ) 

52 await self.create_checkpoint(operation, is_sync=True) 

53 

54 logger.debug( 

55 "Wait checkpoint created for id: %s, name: %s, will suspend", 

56 self.operation_identifier.operation_id, 

57 self.operation_identifier.name, 

58 ) 

59 

60 return await self.execute() 

61 

62 async def replay(self, operation: Operation) -> None: 

63 """Replay an existing wait operation from its checkpoint.""" 

64 if operation.status is OperationStatus.SUCCEEDED: 

65 logger.debug( 

66 "Wait already completed, skipping wait for id: %s, name: %s", 

67 self.operation_identifier.operation_id, 

68 self.operation_identifier.name, 

69 ) 

70 return None 

71 

72 await self.execute() 

73 

74 async def execute(self, operation: Operation | None = None) -> None: 

75 """Execute wait by suspending. 

76 

77 Wait operations 'execute' by suspending execution until the timer completes. 

78 This method never returns normally - it always suspends. 

79 

80 Raises: 

81 SuspendExecution: Always suspends to wait for timer completion 

82 """ 

83 msg: str = f"Wait for {self.seconds} seconds" 

84 suspend_with_optional_resume_delay(msg, self.seconds) # throws suspend 

85 

86 

87def wait(duration: Duration, *, name: str | None = None) -> asyncio.Task[None]: 

88 """Compatibility import for the canonical operation-layer helper.""" 

89 from .._operation.wait import wait as operation_wait 

90 

91 return operation_wait(duration, name=name) 

92 

93 

94async def _wait( 

95 *, 

96 seconds: int, 

97 context: DurableContext, 

98 operation_identifier: OperationIdentifier, 

99) -> None: 

100 executor: WaitOperationExecutor = WaitOperationExecutor( 

101 seconds=seconds, 

102 state=context.execution_state, 

103 operation_identifier=operation_identifier, 

104 ) 

105 await executor.process()