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

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

2 

3from __future__ import annotations 

4 

5import asyncio 

6import logging 

7from typing import TYPE_CHECKING 

8 

9from .base import OperationExecutor 

10from .child import get_durable_context 

11from ..config import Duration, duration_to_seconds 

12from ..exceptions import ValidationError, suspend_with_optional_resume_delay 

13from ..models import ( 

14 Operation, 

15 OperationIdentifier, 

16 OperationStatus, 

17 OperationSubType, 

18 OperationUpdate, 

19 WaitOptions, 

20) 

21from ..task import create_eager_task 

22 

23if TYPE_CHECKING: 

24 from .child import DurableContext 

25 from ..state import ExecutionState 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30class WaitOperationExecutor(OperationExecutor[None]): 

31 """Executor for wait operations.""" 

32 

33 def __init__( 

34 self, 

35 seconds: int, 

36 state: ExecutionState, 

37 operation_identifier: OperationIdentifier, 

38 ): 

39 """Initialize the wait operation executor. 

40 

41 Args: 

42 seconds: Number of seconds to wait 

43 state: The execution state 

44 operation_identifier: The operation identifier 

45 """ 

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

47 self.seconds = seconds 

48 

49 async def start(self) -> None: 

50 """Start a new wait operation.""" 

51 operation: OperationUpdate = OperationUpdate.create_wait_start( 

52 identifier=self.operation_identifier, 

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

54 ) 

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

56 

57 logger.debug( 

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

59 self.operation_identifier.operation_id, 

60 self.operation_identifier.name, 

61 ) 

62 

63 return await self.execute() 

64 

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

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

67 if operation.status is OperationStatus.SUCCEEDED: 

68 logger.debug( 

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

70 self.operation_identifier.operation_id, 

71 self.operation_identifier.name, 

72 ) 

73 return None 

74 

75 await self.execute() 

76 

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

78 """Execute wait by suspending. 

79 

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

81 This method never returns normally - it always suspends. 

82 

83 Raises: 

84 SuspendExecution: Always suspends to wait for timer completion 

85 """ 

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

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

88 

89 

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

91 """Suspend the durable execution for at least the given duration. 

92 

93 Args: 

94 duration: Seconds or timedelta to pause. Must be at least one second. 

95 name: Optional operation name shown in execution history. 

96 """ 

97 context = get_durable_context("wait") 

98 seconds = duration_to_seconds(duration) 

99 if seconds < 1: 

100 msg = "duration must be at least 1 second" 

101 raise ValidationError(msg) 

102 

103 with context._replay_aware(): 

104 operation_id = context.step_counter.create_step_id() 

105 operation_identifier = OperationIdentifier( 

106 operation_id=operation_id, 

107 sub_type=OperationSubType.WAIT, 

108 parent_id=context.parent_id, 

109 name=name, 

110 ) 

111 

112 return create_eager_task( 

113 lambda: _wait( 

114 seconds=seconds, 

115 context=context, 

116 operation_identifier=operation_identifier, 

117 ), 

118 ) 

119 

120 

121async def _wait( 

122 *, 

123 seconds: int, 

124 context: DurableContext, 

125 operation_identifier: OperationIdentifier, 

126) -> None: 

127 executor: WaitOperationExecutor = WaitOperationExecutor( 

128 seconds=seconds, 

129 state=context.execution_state, 

130 operation_identifier=operation_identifier, 

131 ) 

132 await executor.process()