Coverage for async_durable_execution/_operation/recurse.py: 96%

1"""Recursive self-invocation built on the backend-supported invoke operation.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from typing import TypeVar, cast 

7 

8from .._core import ( 

9 RECURSIVE_LEVEL_INPUT_FIELD, 

10 DurableContext, 

11 OperationSubType, 

12 SerDes, 

13 ValidationError, 

14 get_durable_context, 

15) 

16from ..extension import get_extension_context 

17 

18 

19P = TypeVar("P") 

20R = TypeVar("R") 

21 

22 

23def _is_qualified_function_arn(function_name: str) -> bool: 

24 """Return whether a Lambda function ARN includes a qualifier.""" 

25 parts = function_name.split(":") 

26 return len(parts) >= 8 and parts[0] == "arn" and parts[5] == "function" 

27 

28 

29def _is_qualified_function_name(function_name: str) -> bool: 

30 """Return whether a short Lambda function name includes a qualifier.""" 

31 return ":" in function_name and not function_name.startswith("arn:") 

32 

33 

34def _append_qualifier(function_name: str, qualifier: str | None) -> str: 

35 """Append a Lambda qualifier when one is available and needed.""" 

36 if not qualifier: 

37 return function_name 

38 if _is_qualified_function_arn(function_name) or _is_qualified_function_name( 

39 function_name 

40 ): 

41 return function_name 

42 return f"{function_name}:{qualifier}" 

43 

44 

45def _resolve_recursive_function_name( 

46 context: DurableContext, 

47 explicit_function_name: str | None, 

48) -> str: 

49 if explicit_function_name: 

50 return explicit_function_name 

51 

52 lambda_context = context.lambda_context 

53 if lambda_context is None: 53 ↛ 54line 53 didn't jump to line 54 because the condition on line 53 was never true

54 msg = "recurse requires a Lambda context or explicit function_name." 

55 raise RuntimeError(msg) 

56 

57 function_version = getattr(lambda_context, "function_version", None) 

58 invoked_function_arn = getattr(lambda_context, "invoked_function_arn", None) 

59 if invoked_function_arn: 

60 return _append_qualifier(invoked_function_arn, function_version) 

61 

62 context_function_name = getattr(lambda_context, "function_name", None) 

63 if context_function_name: 

64 return _append_qualifier(context_function_name, function_version) 

65 

66 msg = "recurse could not determine the current Lambda function name." 

67 raise RuntimeError(msg) 

68 

69 

70def _validate_recursive_payload(context: DurableContext, payload: P) -> None: 

71 current_input = context.execution_state.get_input_event() 

72 if payload == current_input: 

73 msg = "recurse payload must differ from the current execution input." 

74 raise ValidationError(msg) 

75 

76 

77def _add_recursive_level(context: DurableContext, payload: P) -> P: 

78 if not isinstance(payload, dict): 

79 msg = "recurse payload must be a dict when with_recursive_level is enabled." 

80 raise ValidationError(msg) 

81 

82 payload_with_level = payload.copy() 

83 payload_with_level[RECURSIVE_LEVEL_INPUT_FIELD] = context.recursive_level + 1 

84 return cast("P", payload_with_level) 

85 

86 

87def recurse( 

88 payload: P, 

89 *, 

90 name: str | None = None, 

91 function_name: str | None = None, 

92 with_recursive_level: bool = False, 

93 serdes_payload: SerDes[P] | None = None, 

94 serdes_result: SerDes[R] | None = None, 

95 tenant_id: str | None = None, 

96) -> asyncio.Task[R]: 

97 """Invoke the current durable Lambda function and wait for its result. 

98 

99 This is a convenience wrapper around :func:`invoke` for recursive workflows 

100 such as divide-and-conquer algorithms. Each recursive call is a separate 

101 durable execution, so the current execution records a chained invoke instead 

102 of growing a Python call stack. 

103 

104 Args: 

105 payload: Payload to send to the recursive invocation. 

106 name: Optional durable operation name. 

107 function_name: Optional qualified Lambda function name or ARN. When omitted, 

108 the current Lambda context is used. 

109 with_recursive_level: When true, copy the payload and set 

110 ``__recursive_level`` to the current context level plus one. 

111 serdes_payload: Optional serializer for the invocation payload. 

112 serdes_result: Optional deserializer for the invocation result. 

113 tenant_id: Optional tenant identifier. Defaults to the current Lambda context 

114 tenant id when present. 

115 """ 

116 context = get_durable_context() 

117 recursive_payload = ( 

118 _add_recursive_level(context, payload) if with_recursive_level else payload 

119 ) 

120 _validate_recursive_payload(context, recursive_payload) 

121 target_function_name = _resolve_recursive_function_name(context, function_name) 

122 recursive_tenant_id = tenant_id 

123 if recursive_tenant_id is None and context.lambda_context is not None: 

124 recursive_tenant_id = getattr(context.lambda_context, "tenant_id", None) 

125 

126 return ( 

127 get_extension_context() 

128 ._reserve_sdk_operation(name) # noqa: SLF001 

129 ._run_invoke( # noqa: SLF001 

130 target_function_name, 

131 recursive_payload, 

132 sub_type=OperationSubType.CHAINED_INVOKE, 

133 serdes_payload=serdes_payload, 

134 serdes_result=serdes_result, 

135 tenant_id=recursive_tenant_id, 

136 ) 

137 )