Coverage for async_durable_execution/_core/client.py: 100%

1from __future__ import annotations 

2 

3import asyncio 

4import inspect 

5import importlib.util 

6import logging 

7from collections.abc import Awaitable, Mapping 

8from typing import Any, Protocol, cast 

9 

10from botocore.config import Config 

11from botocore.session import Session 

12 

13from ..__about__ import __version__ 

14from .aws_http import create_botocore_http_client, create_httpx_client 

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(self, **kwargs: Any) -> Mapping[str, Any]: ... 

29 

30 def get_durable_execution_state(self, **kwargs: Any) -> Mapping[str, Any]: ... 

31 

32 

33class AsyncLambdaApiClient(Protocol): 

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

35 

36 def checkpoint_durable_execution( 

37 self, **kwargs: Any 

38 ) -> Awaitable[Mapping[str, Any]]: ... 

39 

40 def get_durable_execution_state( 

41 self, **kwargs: Any 

42 ) -> Awaitable[Mapping[str, Any]]: ... 

43 

44 

45class DurableServiceClient(Protocol): 

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

47 

48 async def checkpoint( 

49 self, 

50 durable_execution_arn: str, 

51 checkpoint_token: str, 

52 updates: list[OperationUpdate], 

53 client_token: str | None, 

54 ) -> CheckpointOutput: ... 

55 

56 async def get_execution_state( 

57 self, 

58 durable_execution_arn: str, 

59 checkpoint_token: str, 

60 next_marker: str, 

61 max_items: int = 1000, 

62 ) -> StateOutput: ... 

63 

64 

65def _create_client_config() -> Config: 

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

67 return Config( 

68 connect_timeout=5, 

69 read_timeout=50, 

70 user_agent_extra=user_agent, 

71 ) 

72 

73 

74def httpx_is_installed() -> bool: 

75 """Return whether the optional HTTPX transport is available.""" 

76 return importlib.util.find_spec("httpx") is not None 

77 

78 

79# Backward-compatible internal alias for integrations that imported the old name. 

80aioboto_is_installed = httpx_is_installed 

81 

82 

83def create_default_sync_client( 

84 *, 

85 session: Session | None = None, 

86 endpoint_url: str | None = None, 

87 region_name: str | None = None, 

88 config: Config | None = None, 

89) -> LambdaApiClient: 

90 """Create a model-free Lambda client using botocore's HTTP transport.""" 

91 return create_botocore_http_client( 

92 session=session, 

93 endpoint_url=endpoint_url, 

94 region_name=region_name, 

95 config=config or _create_client_config(), 

96 ) 

97 

98 

99def create_default_async_client( 

100 *, 

101 session: Session | None = None, 

102 endpoint_url: str | None = None, 

103 region_name: str | None = None, 

104 config: Config | None = None, 

105) -> AsyncLambdaApiClient: 

106 """Create a model-free Lambda client using HTTPX.""" 

107 return create_httpx_client( 

108 session=session, 

109 endpoint_url=endpoint_url, 

110 region_name=region_name, 

111 config=config or _create_client_config(), 

112 ) 

113 

114 

115def create_default_client() -> LambdaApiClient | AsyncLambdaApiClient: 

116 """Create the default Lambda client, preferring async when HTTPX is installed.""" 

117 if httpx_is_installed(): 

118 return create_default_async_client() 

119 return create_default_sync_client() 

120 

121 

122def lambda_api_client_is_async( 

123 client: LambdaApiClient | AsyncLambdaApiClient, 

124) -> bool: 

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

126 return inspect.iscoroutinefunction(client.checkpoint_durable_execution) 

127 

128 

129def create_default_service_client( 

130 client: LambdaApiClient | AsyncLambdaApiClient | None = None, 

131) -> DurableServiceClient: 

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

133 lambda_client = client or create_default_client() 

134 if lambda_api_client_is_async(lambda_client): 

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

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

137 

138 

139class ThreadedSyncLambdaClient(DurableServiceClient): 

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

141 

142 _cached_boto_client: LambdaApiClient | None = None 

143 

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

145 self.client = client or create_default_sync_client() 

146 

147 async def checkpoint( 

148 self, 

149 durable_execution_arn: str, 

150 checkpoint_token: str | None, 

151 updates: list[OperationUpdate], 

152 client_token: str | None, 

153 ) -> CheckpointOutput: 

154 if not checkpoint_token: 

155 msg = "Cannot checkpoint without a checkpoint token." 

156 raise CheckpointError(msg) 

157 

158 try: 

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

160 if client_token is not None: 

161 optional_params["ClientToken"] = client_token 

162 

163 result = await asyncio.to_thread( 

164 self.client.checkpoint_durable_execution, 

165 DurableExecutionArn=durable_execution_arn, 

166 CheckpointToken=checkpoint_token, 

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

168 **optional_params, 

169 ) 

170 

171 return CheckpointOutput.from_dict(result) 

172 except Exception as e: 

173 checkpoint_error = CheckpointError.from_exception(e) 

174 logger.exception( 

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

176 ) 

177 raise checkpoint_error from None 

178 

179 async def get_execution_state( 

180 self, 

181 durable_execution_arn: str, 

182 checkpoint_token: str | None, 

183 next_marker: str, 

184 max_items: int = 1000, 

185 ) -> StateOutput: 

186 if not checkpoint_token: 

187 msg = "Cannot get execution state without a checkpoint token." 

188 raise GetExecutionStateError(msg) 

189 

190 try: 

191 result = await asyncio.to_thread( 

192 self.client.get_durable_execution_state, 

193 DurableExecutionArn=durable_execution_arn, 

194 CheckpointToken=checkpoint_token, 

195 Marker=next_marker, 

196 MaxItems=max_items, 

197 ) 

198 return StateOutput.from_dict(result) 

199 except Exception as e: 

200 error = GetExecutionStateError.from_exception(e) 

201 logger.exception( 

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

203 ) 

204 raise error from None 

205 

206 

207class AsyncLambdaClient(DurableServiceClient): 

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

209 

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

211 self.client = client 

212 

213 async def checkpoint( 

214 self, 

215 durable_execution_arn: str, 

216 checkpoint_token: str | None, 

217 updates: list[OperationUpdate], 

218 client_token: str | None, 

219 ) -> CheckpointOutput: 

220 if not checkpoint_token: 

221 msg = "Cannot checkpoint without a checkpoint token." 

222 raise CheckpointError(msg) 

223 

224 try: 

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

226 if client_token is not None: 

227 optional_params["ClientToken"] = client_token 

228 

229 result = await self.client.checkpoint_durable_execution( 

230 DurableExecutionArn=durable_execution_arn, 

231 CheckpointToken=checkpoint_token, 

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

233 **optional_params, 

234 ) 

235 

236 return CheckpointOutput.from_dict(result) 

237 except Exception as e: 

238 checkpoint_error = CheckpointError.from_exception(e) 

239 logger.exception( 

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

241 ) 

242 raise checkpoint_error from None 

243 

244 async def get_execution_state( 

245 self, 

246 durable_execution_arn: str, 

247 checkpoint_token: str | None, 

248 next_marker: str, 

249 max_items: int = 1000, 

250 ) -> StateOutput: 

251 if not checkpoint_token: 

252 msg = "Cannot get execution state without a checkpoint token." 

253 raise GetExecutionStateError(msg) 

254 

255 try: 

256 result = await self.client.get_durable_execution_state( 

257 DurableExecutionArn=durable_execution_arn, 

258 CheckpointToken=checkpoint_token, 

259 Marker=next_marker, 

260 MaxItems=max_items, 

261 ) 

262 return StateOutput.from_dict(result) 

263 except Exception as e: 

264 error = GetExecutionStateError.from_exception(e) 

265 logger.exception( 

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

267 ) 

268 raise error from None 

269 

270 async def aclose(self) -> None: 

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

272 if close is None: 

273 return 

274 await close() 

275 

276 

277__all__ = [ 

278 "AsyncLambdaClient", 

279 "DurableServiceClient", 

280 "ThreadedSyncLambdaClient", 

281 "create_default_async_client", 

282 "create_default_client", 

283 "create_default_service_client", 

284 "create_default_sync_client", 

285 "lambda_api_client_is_async", 

286]