Coverage for source/aws/aws_handler.py: 74%

1# aws/aws_handler.py 

2 

3import os 

4import boto3 

5import io 

6 

7class AWSHandler: 

8 """ 

9 Responsible for handling communication with Amazon AWS services. 

10 """ 

11 

12 def __init__(self, role_name: str, region_name: str = "eu-central-1") -> None: 

13 """ 

14 Class constructor. Before calling it AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY 

15 and ACCOUNT_ID should be available as environmental variables. 

16 

17 Parameters: 

18 role_arn (str): Assumed role identification string. 

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

20 

21 Raises: 

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

23 """ 

24 

25 AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID') 

26 AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY') 

27 ACCOUNT_ID = os.getenv('ACCOUNT_ID') 

28 if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY or not ACCOUNT_ID: 

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

30 

31 session = boto3.Session(aws_access_key_id = AWS_ACCESS_KEY_ID, 

32 aws_secret_access_key = AWS_SECRET_ACCESS_KEY) 

33 role_arn = f'arn:aws:iam::{ACCOUNT_ID}:role/{role_name}' 

34 

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

36 RoleSessionName = 'S3_bucket_user_session') 

37 credentials = assumed_role['Credentials'] 

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

39 aws_secret_access_key = credentials['SecretAccessKey'], 

40 aws_session_token = credentials['SessionToken'], 

41 region_name = region_name) 

42 

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

44 """ 

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

46 

47 Parameters: 

48 bucket_name (str): String denoting bucket name. 

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

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

51 If left unspecified, name does not change. 

52 

53 Raises: 

54 RuntimeError: If approached problem during file uploading. 

55 """ 

56 

57 if desired_name == "": 

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

59 try: 

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

61 except Exception as e: 

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

63 

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

65 """ 

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

67 

68 Parameters: 

69 bucket_name (str): String denoting bucket name. 

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

71 written to bucket. 

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

73 

74 Raises: 

75 RuntimeError: If approached problem during file uploading. 

76 """ 

77 

78 try: 

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

80 Body = buffer.getvalue()) 

81 except Exception as e: 

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

83 

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

85 """ 

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

87 

88 Parameters: 

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

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

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

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

93 with the original filename. 

94 

95 Raises: 

96 RuntimeError: If the download operation fails. 

97 """ 

98 

99 if desired_path == "": 

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

101 try: 

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

103 except Exception as e: 

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