Coverage for source/utils/aws_handler.py: 76%

1# utils/aws_handler.py 

2 

3# global imports 

4import boto3 

5import io 

6import os 

7 

8# local imports 

9from source.utils import SingletonMeta 

10 

11class AWSHandler(metaclass = SingletonMeta): 

12 """ 

13 Responsible for handling communication with Amazon AWS services. 

14 """ 

15 

16 # local constants 

17 __DEFAULT_REGION = "eu-central-1" 

18 

19 def __init__(self, region_name: str = __DEFAULT_REGION) -> None: 

20 """ 

21 Class constructor. Before calling it AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY 

22 and ACCOUNT_ID should be available as environmental variables. 

23 

24 Parameters: 

25 region_name (str): Region name to connect to. 

26 

27 Raises: 

28 RuntimeError: If AWS credentials or account ID are not defined. 

29 """ 

30 

31 AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID') 

32 AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY') 

33 ACCOUNT_ID = os.getenv('ACCOUNT_ID') 

34 ROLE_NAME = os.getenv('ROLE_NAME') 

35 if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY or not ACCOUNT_ID \ 

36 or not ROLE_NAME: 

37 raise RuntimeError("AWS credentials or account ID not found in environment variables!") 

38 

39 session = boto3.Session(aws_access_key_id = AWS_ACCESS_KEY_ID, 

40 aws_secret_access_key = AWS_SECRET_ACCESS_KEY) 

41 role_arn = f'arn:aws:iam::{ACCOUNT_ID}:role/{ROLE_NAME}' 

42 

43 assumed_role = session.client('sts').assume_role(RoleArn = role_arn, 

44 RoleSessionName = 'S3_bucket_user_session') 

45 credentials = assumed_role['Credentials'] 

46 self.aws_s3_resource = boto3.client('s3', aws_access_key_id = credentials['AccessKeyId'], 

47 aws_secret_access_key = credentials['SecretAccessKey'], 

48 aws_session_token = credentials['SessionToken'], 

49 region_name = region_name) 

50 

51 def upload_file_to_s3(self, bucket_name: str, file_path: str, desired_name: str = "") -> None: 

52 """ 

53 Attempts to upload local file specified by path to S3 Amazon bucket. 

54 

55 Parameters: 

56 bucket_name (str): String denoting bucket name. 

57 file_path (str): String representing file to the path that should be uploaded. 

58 desired_name (str): Desired name to be given to the file after being uploaded. 

59 If left unspecified, name does not change. 

60 

61 Raises: 

62 RuntimeError: If approached problem during file uploading. 

63 """ 

64 

65 if desired_name == "": 

66 desired_name = file_path.split('/')[-1] 

67 try: 

68 self.aws_s3_resource.upload_file(file_path, bucket_name, desired_name) 

69 except Exception as e: 

70 raise RuntimeError(f"Did not managed to upload file! Original error: {e}") 

71 

72 def upload_buffer_to_s3(self, bucket_name: str, buffer: io.StringIO, desired_name: str) -> None: 

73 """ 

74 Attempts to upload buffer as file body directly to S3 Amazon bucket. 

75 

76 Parameters: 

77 bucket_name (str): String denoting bucket name. 

78 buffer (io.StringIO): Buffer containing data that should be directly 

79 written to bucket. 

80 desired_name (str): Desired name to be given to the file after being uploaded. 

81 

82 Raises: 

83 RuntimeError: If approached problem during file uploading. 

84 """ 

85 

86 try: 

87 self.aws_s3_resource.put_object(Bucket = bucket_name, Key = desired_name, 

88 Body = buffer.getvalue()) 

89 except Exception as e: 

90 raise RuntimeError(f"Did not managed to upload file! Original error: {e}") 

91 

92 def download_file_from_s3(self, bucket_name: str, file_name: str, desired_path: str = "") -> None: 

93 """ 

94 Downloads a file from an S3 bucket to a local path. 

95 

96 Parameters: 

97 bucket_name (str): The name of the S3 bucket. 

98 file_name (str): The key/path of the file in the S3 bucket. 

99 desired_path (str, optional): The local path where the file will be saved. 

100 If not provided, the file will be downloaded to the current working directory 

101 with the original filename. 

102 

103 Raises: 

104 RuntimeError: If the download operation fails. 

105 """ 

106 

107 if desired_path == "": 

108 desired_path = os.getcwd() + '/' + file_name.split('/')[-1] 

109 try: 

110 self.aws_s3_resource.download_file(bucket_name, file_name, desired_path) 

111 except Exception as e: 

112 raise RuntimeError(f"Did not managed to download file! Original error: {e}")