[feat] #15 Add S3 storage strategy - better handle exceptions and move resolve placeholder to storage strategy

Author: choinek

app/storage_strategies/aws_s3.py

@@ -1,29 +1,37 @@
 import boto3
-import os
-from string import Template
-from botocore.exceptions import NoCredentialsError, ClientError
+from botocore.exceptions import EndpointConnectionError, ClientError
 from storage_strategies.storage_strategy import StorageStrategy
 
 class AWSS3StorageStrategy(StorageStrategy):
     def __init__(self, context):
         super().__init__(context)
 
-        self.bucket_name = self._resolve_placeholder(context['settings'].get('bucket_name'))
-        self.region = self._resolve_placeholder(context['settings'].get('region'))
-        self.access_key = self._resolve_placeholder(context['settings'].get('access_key'))
-        self.secret_access_key = self._resolve_placeholder(context['settings'].get('secret_access_key'))
+        self.bucket_name = self.resolve_placeholder(context['settings'].get('bucket_name'))
+        self.region = self.resolve_placeholder(context['settings'].get('region'))
+        self.access_key = self.resolve_placeholder(context['settings'].get('access_key'))
+        self.secret_access_key = self.resolve_placeholder(context['settings'].get('secret_access_key'))
 
-        self.s3_client = boto3.client(
-            's3',
-            aws_access_key_id=self.access_key,
-            aws_secret_access_key=self.secret_access_key,
-            region_name=self.region
-        )
-
-    def _resolve_placeholder(self, value):
-        if not value:
-            return None
-        return Template(value).substitute(os.environ)
+        try:
+            self.s3_client = boto3.client(
+                's3',
+                aws_access_key_id=self.access_key,
+                aws_secret_access_key=self.secret_access_key,
+                region_name=self.region
+            )
+            self.s3_client.head_bucket(Bucket=self.bucket_name)
+        except EndpointConnectionError as e:
+            raise RuntimeError(
+                f"{str(e)}\n"
+                "Check your AWS_REGION and AWS_S3_BUCKET_NAME environment variables."
+            ) from e
+        except ClientError as e:
+            error_code = e.response.get('Error', {}).get('Code', 'Unknown')
+            if error_code in ('400', '403'):
+                raise RuntimeError(
+                    f"{str(e)}\n"
+                    "Error: Please check your AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY."
+                ) from e
+            raise
 
     def save(self, file_name, dest_file_name, content):
         formatted_file_name = self.format_file_name(file_name, dest_file_name)
@@ -34,31 +42,40 @@ def save(self, file_name, dest_file_name, content):
                 Key=formatted_file_name,
                 Body=content.encode('utf-8')
             )
-            print(f"File {formatted_file_name} saved to S3 bucket {self.bucket_name}.")
-        except NoCredentialsError:
-            print("AWS credentials are missing or invalid.")
         except ClientError as e:
-            print(f"Failed to upload {file_name}: {e}")
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error saving file '{file_name}' as '{formatted_file_name}' to bucket '{self.bucket_name}'."
+            ) from e
 
     def load(self, file_name):
         try:
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=file_name)
             return response['Body'].read().decode('utf-8')
         except ClientError as e:
-            print(f"Failed to download {file_name}: {e}")
-            return None
+            error_code = e.response['Error']['Code']
+            if error_code == 'NoSuchKey':
+                return None
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error loading file '{file_name}' from bucket '{self.bucket_name}'."
+            ) from e
 
     def list(self):
         try:
             response = self.s3_client.list_objects_v2(Bucket=self.bucket_name)
             return [item['Key'] for item in response.get('Contents', [])]
         except ClientError as e:
-            print(f"Failed to list files: {e}")
-            return []
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error listing objects in bucket '{self.bucket_name}'."
+            ) from e
 
     def delete(self, file_name):
         try:
             self.s3_client.delete_object(Bucket=self.bucket_name, Key=file_name)
-            print(f"File {file_name} deleted from S3 bucket {self.bucket_name}.")
         except ClientError as e:
-            print(f"Failed to delete {file_name}: {e}")
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error deleting file '{file_name}' from bucket '{self.bucket_name}'."
+            ) from e

app/storage_strategies/storage_strategy.py

@@ -2,6 +2,7 @@
 import yaml
 from datetime import datetime
 from pathlib import Path
+from string import Template
 
 class StorageStrategy:
     def __init__(self, context):
@@ -18,7 +19,7 @@ def list(self):
 
     def delete(self, file_name):
         raise NotImplementedError("Subclasses must implement this method")
-    
+
     def format_file_name(self, file_name, format_string):
         return format_string.format(file_fullname=file_name,  # file_name with path
                                      file_name=Path(file_name).stem,  # file_name without path
@@ -28,4 +29,15 @@ def format_file_name(self, file_name, format_string):
                                      dd=datetime.now().strftime('%d'),
                                      HH=datetime.now().strftime('%H'),
                                      MM=datetime.now().strftime('%M'),
-                                     SS=datetime.now().strftime('%S'))
+                                     SS=datetime.now().strftime('%S'))
+
+    def resolve_placeholder(self, value, default=None):
+        if not value:
+            return default
+        try:
+            return Template(value).substitute(os.environ)
+        except KeyError as e:
+            if default:
+                return default
+            else:
+                raise ValueError(f"Environment variable '{e.args[0]}' is missing, and no default value is provided.")