From 5623734276d93f547eeb3950230cbd362e9547f4 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Thu, 10 Jul 2025 15:34:52 +0530 Subject: [PATCH] feat(storage): enhance save_file method to accept storage class parameter --- application/api/user/routes.py | 2 +- application/storage/base.py | 3 ++- application/storage/s3.py | 12 ++++++++++-- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/application/api/user/routes.py b/application/api/user/routes.py index c2f89761..11594161 100644 --- a/application/api/user/routes.py +++ b/application/api/user/routes.py @@ -168,7 +168,7 @@ def handle_image_upload( filename = secure_filename(file.filename) upload_path = f"{settings.UPLOAD_FOLDER.rstrip('/')}/{user}/{base_path.rstrip('/')}/{uuid.uuid4()}_{filename}" try: - storage.save_file(file, upload_path) + storage.save_file(file, upload_path, storage_class="STANDARD") image_url = upload_path except Exception as e: current_app.logger.error(f"Error uploading image: {e}") diff --git a/application/storage/base.py b/application/storage/base.py index 273e7761..07f33c7b 100644 --- a/application/storage/base.py +++ b/application/storage/base.py @@ -1,4 +1,5 @@ """Base storage class for file system abstraction.""" + from abc import ABC, abstractmethod from typing import BinaryIO, List, Callable @@ -7,7 +8,7 @@ class BaseStorage(ABC): """Abstract base class for storage implementations.""" @abstractmethod - def save_file(self, file_data: BinaryIO, path: str) -> dict: + def save_file(self, file_data: BinaryIO, path: str, **kwargs) -> dict: """ Save a file to storage. diff --git a/application/storage/s3.py b/application/storage/s3.py index 7a52dd1c..1babb843 100644 --- a/application/storage/s3.py +++ b/application/storage/s3.py @@ -38,9 +38,17 @@ class S3Storage(BaseStorage): region_name=region_name, ) - def save_file(self, file_data: BinaryIO, path: str) -> dict: + def save_file( + self, + file_data: BinaryIO, + path: str, + storage_class: str = "INTELLIGENT_TIERING", + **kwargs, + ) -> dict: """Save a file to S3 storage.""" - self.s3.upload_fileobj(file_data, self.bucket_name, path) + self.s3.upload_fileobj( + file_data, self.bucket_name, path, ExtraArgs={"StorageClass": storage_class} + ) region = getattr(settings, "SAGEMAKER_REGION", None)