Coverage for src/extratools_cloud/aws/helpers.py: 0%

1from collections.abc import Callable 

2from os import getenv 

3from typing import Any 

4 

5import boto3 

6from boto3.resources.base import ServiceResource 

7from botocore.client import BaseClient 

8from botocore.config import Config 

9from botocore.exceptions import ClientError 

10 

11LOCAL_STAGE: str = "local" 

12 

13STAGE: str = getenv("STAGE", LOCAL_STAGE) 

14 

15LOCALSTACK_ENDPOINT: str = "http://localhost:4566" 

16 

17 

18def get_client(service: str, *, config: Config | None = None) -> BaseClient: 

19 return boto3.client( 

20 service, 

21 endpoint_url=( 

22 LOCALSTACK_ENDPOINT if STAGE == LOCAL_STAGE 

23 else None 

24 ), 

25 config=config, 

26 ) 

27 

28 

29def get_service_resource(service: str) -> ServiceResource: 

30 return boto3.resource( 

31 service, 

32 endpoint_url=( 

33 LOCALSTACK_ENDPOINT if STAGE == LOCAL_STAGE 

34 else None 

35 ), 

36 ) 

37 

38 

39class ClientErrorHandler: 

40 """ 

41 Based on https://boto3.amazonaws.com/v1/documentation/api/latest/guide/error-handling.html 

42 """ 

43 

44 def __init__( 

45 self, 

46 error_code: str, 

47 exception_class: type[Exception], 

48 ) -> None: 

49 self.__error_code = error_code 

50 self.__exception_class = exception_class 

51 

52 def __call__(self, f: Callable[..., Any]) -> Any: 

53 def wrapper(*args: Any, **kwargs: Any) -> Any: 

54 try: 

55 return f(*args, **kwargs) 

56 except ClientError as e: 

57 error = e.response["Error"] 

58 # For SQS, somehow error code is found in 

59 # both `Code` and `QueryErrorCode` with different values 

60 # (`AWS.SimpleQueueService.NonExistentQueue` vs `QueueDoesNotExist`). 

61 # The value in `QueryErrorCode` seems to match public API doc. 

62 # https://github.com/aws/aws-sdk/issues/105 

63 if error.get("QueryErrorCode", error["Code"]) == self.__error_code: 

64 raise self.__exception_class from e 

65 

66 raise 

67 

68 return wrapper