Coverage for src/tagmania/iac_tools/util.py: 50%

1"""AWS Utility Functions for Tagmania. 

2 

3This module provides utility functions for working with various AWS services 

4beyond the core EC2 functionality. Currently focuses on AWS Lambda functions 

5but can be extended for other services as needed. 

6 

7The utilities include helper functions for retrieving ARNs, managing tags, 

8and working with AWS clients in a consistent manner across the application. 

9 

10Example: 

11 Lambda function utilities: 

12 

13 ```python 

14 # Get Lambda function ARN 

15 arn = get_lambda_arn('my-function') 

16 

17 # Tag multiple functions 

18 arns = [get_lambda_arn('func1'), get_lambda_arn('func2')] 

19 tags = {'Environment': 'production', 'Team': 'backend'} 

20 tag_lambda_functions(arns, tags) 

21 ``` 

22""" 

23 

24import os 

25from functools import lru_cache 

26from typing import Any 

27 

28import boto3 

29 

30 

31@lru_cache # cache the client the first time we ask for it 

32def _lambda_client() -> Any: 

33 """Get cached AWS Lambda client with fallback region handling. 

34 

35 Returns: 

36 boto3.client: Configured Lambda client with region fallback 

37 

38 Note: 

39 Uses LRU cache to avoid creating multiple clients. Falls back to 

40 us-east-1 if AWS_REGION environment variable is not set. 

41 """ 

42 # Fall back to env-var or default region so boto3 never errors 

43 return boto3.client("lambda", region_name=os.getenv("AWS_REGION", "us-east-1")) 

44 

45 

46def get_lambda_arn(function_name: str) -> str: 

47 """Retrieve the ARN for a Lambda function. 

48 

49 Args: 

50 function_name: Name of the Lambda function 

51 

52 Returns: 

53 str: Full ARN of the Lambda function 

54 

55 Raises: 

56 ClientError: If function doesn't exist or access is denied 

57 

58 Example: 

59 ```python 

60 arn = get_lambda_arn('my-processing-function') 

61 print(f"Function ARN: {arn}") 

62 ``` 

63 """ 

64 response = _lambda_client().get_function(FunctionName=function_name) 

65 return str(response["Configuration"]["FunctionArn"]) 

66 

67 

68def tag_lambda_functions(function_arns: list[str], tags: dict[str, str]) -> None: 

69 """Apply tags to multiple Lambda functions. 

70 

71 Args: 

72 function_arns: List of Lambda function ARNs 

73 tags: Dictionary of tag key-value pairs 

74 

75 Example: 

76 ```python 

77 arns = [get_lambda_arn('func1'), get_lambda_arn('func2')] 

78 tags = {'Environment': 'production', 'Owner': 'team-backend'} 

79 tag_lambda_functions(arns, tags) 

80 ``` 

81 """ 

82 client = _lambda_client() 

83 for arn in function_arns: 

84 client.tag_resource(Resource=arn, Tags=tags) 

85 

86 

87def untag_lambda_functions(function_arns: list[str], keys: list[str]) -> None: 

88 """Remove tags from multiple Lambda functions. 

89 

90 Args: 

91 function_arns: List of Lambda function ARNs 

92 keys: List of tag keys to remove 

93 

94 Example: 

95 ```python 

96 arns = [get_lambda_arn('func1'), get_lambda_arn('func2')] 

97 keys = ['Environment', 'TempTag'] 

98 untag_lambda_functions(arns, keys) 

99 ``` 

100 """ 

101 client = _lambda_client() 

102 for arn in function_arns: 

103 client.untag_resource(Resource=arn, TagKeys=keys)