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

1from __future__ import annotations 

2 

3import asyncio 

4import inspect 

5import importlib 

6import importlib.util 

7import logging 

8from collections.abc import Awaitable, Mapping 

9from typing import Any, Protocol, cast 

10 

11from botocore.config import Config 

12from botocore.session import get_session 

13 

14from .__about__ import __version__ 

15from .exceptions import CheckpointError, GetExecutionStateError 

16from .models import ( 

17 CheckpointOutput, 

18 OperationUpdate, 

19 StateOutput, 

20) 

21 

22logger = logging.getLogger(__name__) 

23 

24 

25class LambdaApiClient(Protocol): 

26 """Minimal Lambda client surface needed by durable execution.""" 

27 

28 def checkpoint_durable_execution( 

29 self, **kwargs: Any 

30 ) -> Mapping[str, Any]: ... # pragma: no cover 

31 

32 def get_durable_execution_state( 

33 self, **kwargs: Any 

34 ) -> Mapping[str, Any]: ... # pragma: no cover 

35 

36 

37class AsyncLambdaApiClient(Protocol): 

38 """Minimal async Lambda client surface needed by durable execution.""" 

39 

40 def checkpoint_durable_execution( 

41 self, **kwargs: Any 

42 ) -> Awaitable[Mapping[str, Any]]: ... # pragma: no cover 

43 

44 def get_durable_execution_state( 

45 self, **kwargs: Any 

46 ) -> Awaitable[Mapping[str, Any]]: ... # pragma: no cover 

47 

48 

49class DurableServiceClient(Protocol): 

50 """Durable Service clients must implement this interface.""" 

51 

52 async def checkpoint( 

53 self, 

54 durable_execution_arn: str, 

55 checkpoint_token: str, 

56 updates: list[OperationUpdate], 

57 client_token: str | None, 

58 ) -> CheckpointOutput: ... # pragma: no cover 

59 

60 async def get_execution_state( 

61 self, 

62 durable_execution_arn: str, 

63 checkpoint_token: str, 

64 next_marker: str, 

65 max_items: int = 1000, 

66 ) -> StateOutput: ... # pragma: no cover 

67 

68 

69def _create_client_config() -> Config: 

70 user_agent = f"durable-execution-sdk-python/{__version__}-async" 

71 return Config( 

72 connect_timeout=5, 

73 read_timeout=50, 

74 user_agent_extra=user_agent, 

75 ) 

76 

77 

78def aioboto_is_installed() -> bool: 

79 """Return whether the optional aioboto dependency is available.""" 

80 return importlib.util.find_spec("aiobotocore") is not None 

81 

82 

83def create_default_sync_client() -> LambdaApiClient: 

84 """Create the default botocore Lambda client used for durable API calls.""" 

85 session = get_session() 

86 return cast( 

87 "LambdaApiClient", 

88 session.create_client("lambda", config=_create_client_config()), 

89 ) 

90 

91 

92def create_default_async_client() -> AsyncLambdaApiClient: 

93 """Create the default aioboto Lambda client used for durable API calls.""" 

94 aiobotocore_session = importlib.import_module("aiobotocore.session") 

95 session = aiobotocore_session.get_session() 

96 return _AiobotocoreLambdaApiClient( 

97 session.create_client("lambda", config=_create_client_config()) 

98 ) 

99 

100 

101def create_default_client() -> LambdaApiClient | AsyncLambdaApiClient: 

102 """Create the default Lambda client, preferring async when aioboto is installed.""" 

103 if aioboto_is_installed(): 

104 return create_default_async_client() 

105 return create_default_sync_client() 

106 

107 

108def lambda_api_client_is_async( 

109 client: LambdaApiClient | AsyncLambdaApiClient, 

110) -> bool: 

111 """Return whether a Lambda API client exposes async durable methods.""" 

112 return inspect.iscoroutinefunction(client.checkpoint_durable_execution) 

113 

114 

115def create_default_service_client( 

116 client: LambdaApiClient | AsyncLambdaApiClient | None = None, 

117) -> DurableServiceClient: 

118 """Create the default durable service client.""" 

119 lambda_client = client or create_default_client() 

120 if lambda_api_client_is_async(lambda_client): 

121 return AsyncLambdaClient(cast("AsyncLambdaApiClient", lambda_client)) 

122 return ThreadedSyncLambdaClient(cast("LambdaApiClient", lambda_client)) 

123 

124 

125class ThreadedSyncLambdaClient(DurableServiceClient): 

126 """Adapt the sync botocore Lambda client to the async service interface.""" 

127 

128 _cached_boto_client: LambdaApiClient | None = None 

129 

130 def __init__(self, client: LambdaApiClient | None) -> None: 

131 self.client = client or create_default_sync_client() 

132 

133 async def checkpoint( 

134 self, 

135 durable_execution_arn: str, 

136 checkpoint_token: str, 

137 updates: list[OperationUpdate], 

138 client_token: str | None, 

139 ) -> CheckpointOutput: 

140 try: 

141 optional_params: dict[str, str] = {} 

142 if client_token is not None: 

143 optional_params["ClientToken"] = client_token 

144 

145 result = await asyncio.to_thread( 

146 self.client.checkpoint_durable_execution, 

147 DurableExecutionArn=durable_execution_arn, 

148 CheckpointToken=checkpoint_token, 

149 Updates=cast("Any", [o.to_dict() for o in updates]), 

150 **optional_params, 

151 ) 

152 

153 return CheckpointOutput.from_dict(result) 

154 except Exception as e: 

155 checkpoint_error = CheckpointError.from_exception(e) 

156 logger.exception( 

157 "Failed to checkpoint.", extra=checkpoint_error.build_logger_extras() 

158 ) 

159 raise checkpoint_error from None 

160 

161 async def get_execution_state( 

162 self, 

163 durable_execution_arn: str, 

164 checkpoint_token: str, 

165 next_marker: str, 

166 max_items: int = 1000, 

167 ) -> StateOutput: 

168 try: 

169 result = await asyncio.to_thread( 

170 self.client.get_durable_execution_state, 

171 DurableExecutionArn=durable_execution_arn, 

172 CheckpointToken=checkpoint_token, 

173 Marker=next_marker, 

174 MaxItems=max_items, 

175 ) 

176 return StateOutput.from_dict(result) 

177 except Exception as e: 

178 error = GetExecutionStateError.from_exception(e) 

179 logger.exception( 

180 "Failed to get execution state.", extra=error.build_logger_extras() 

181 ) 

182 raise error from None 

183 

184 

185class _AiobotocoreLambdaApiClient: 

186 """Lazily enter an aiobotocore Lambda client context for durable API calls.""" 

187 

188 def __init__(self, client_context: Any) -> None: 

189 self._client_context = client_context 

190 self._client: Any | None = None 

191 

192 async def _get_client(self) -> Any: 

193 if self._client is None: 

194 self._client = await self._client_context.__aenter__() 

195 return self._client 

196 

197 async def checkpoint_durable_execution(self, **kwargs: Any) -> Any: 

198 client = await self._get_client() 

199 return await client.checkpoint_durable_execution(**kwargs) 

200 

201 async def get_durable_execution_state(self, **kwargs: Any) -> Any: 

202 client = await self._get_client() 

203 return await client.get_durable_execution_state(**kwargs) 

204 

205 async def aclose(self) -> None: 

206 if self._client is None: 

207 return 

208 await self._client_context.__aexit__(None, None, None) 

209 self._client = None 

210 

211 

212class AsyncLambdaClient(DurableServiceClient): 

213 """Adapt an async aioboto Lambda client to the durable service interface.""" 

214 

215 def __init__(self, client: AsyncLambdaApiClient) -> None: 

216 self.client = client 

217 

218 async def checkpoint( 

219 self, 

220 durable_execution_arn: str, 

221 checkpoint_token: str, 

222 updates: list[OperationUpdate], 

223 client_token: str | None, 

224 ) -> CheckpointOutput: 

225 try: 

226 optional_params: dict[str, str] = {} 

227 if client_token is not None: 

228 optional_params["ClientToken"] = client_token 

229 

230 result = await self.client.checkpoint_durable_execution( 

231 DurableExecutionArn=durable_execution_arn, 

232 CheckpointToken=checkpoint_token, 

233 Updates=cast("Any", [o.to_dict() for o in updates]), 

234 **optional_params, 

235 ) 

236 

237 return CheckpointOutput.from_dict(result) 

238 except Exception as e: 

239 checkpoint_error = CheckpointError.from_exception(e) 

240 logger.exception( 

241 "Failed to checkpoint.", extra=checkpoint_error.build_logger_extras() 

242 ) 

243 raise checkpoint_error from None 

244 

245 async def get_execution_state( 

246 self, 

247 durable_execution_arn: str, 

248 checkpoint_token: str, 

249 next_marker: str, 

250 max_items: int = 1000, 

251 ) -> StateOutput: 

252 try: 

253 result = await self.client.get_durable_execution_state( 

254 DurableExecutionArn=durable_execution_arn, 

255 CheckpointToken=checkpoint_token, 

256 Marker=next_marker, 

257 MaxItems=max_items, 

258 ) 

259 return StateOutput.from_dict(result) 

260 except Exception as e: 

261 error = GetExecutionStateError.from_exception(e) 

262 logger.exception( 

263 "Failed to get execution state.", extra=error.build_logger_extras() 

264 ) 

265 raise error from None 

266 

267 async def aclose(self) -> None: 

268 close = getattr(self.client, "aclose", None) 

269 if close is None: 

270 return 

271 await close() 

272 

273 

274__all__ = [ 

275 "AsyncLambdaClient", 

276 "DurableServiceClient", 

277 "ThreadedSyncLambdaClient", 

278 "create_default_async_client", 

279 "create_default_client", 

280 "create_default_service_client", 

281 "create_default_sync_client", 

282 "lambda_api_client_is_async", 

283]