Split the storage providers into separate classes in preparation for other storage providers like GCS

backend/open_webui/storage/provider.py

@@ -1,40 +1,85 @@
 import os
-import boto3
-from botocore.exceptions import ClientError
 import shutil
+from abc import ABC, abstractmethod
+from typing import BinaryIO, Tuple
 
-
-from typing import BinaryIO, Tuple, Optional, Union
-
-from open_webui.constants import ERROR_MESSAGES
+import boto3
+from botocore.exceptions import ClientError
 from open_webui.config import (
-    STORAGE_PROVIDER,
     S3_ACCESS_KEY_ID,
-    S3_SECRET_ACCESS_KEY,
     S3_BUCKET_NAME,
-    S3_REGION_NAME,
     S3_ENDPOINT_URL,
+    S3_REGION_NAME,
+    S3_SECRET_ACCESS_KEY,
+    STORAGE_PROVIDER,
     UPLOAD_DIR,
 )
+from open_webui.constants import ERROR_MESSAGES
 
 
-import boto3
-from botocore.exceptions import ClientError
-from typing import BinaryIO, Tuple, Optional
+class StorageProvider(ABC):
+    @abstractmethod
+    def get_file(self, file_path: str) -> str:
+        pass
 
+    @abstractmethod
+    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+        pass
 
-class StorageProvider:
-    def __init__(self, provider: Optional[str] = None):
-        self.storage_provider: str = provider or STORAGE_PROVIDER
+    @abstractmethod
+    def delete_all_files(self) -> None:
+        pass
+
+    @abstractmethod
+    def delete_file(self, file_path: str) -> None:
+        pass
 
-        self.s3_client = None
-        self.s3_bucket_name: Optional[str] = None
 
-        if self.storage_provider == "s3":
-            self._initialize_s3()
+class LocalStorageProvider(StorageProvider):
+    @staticmethod
+    def upload_file(file: BinaryIO, filename: str) -> Tuple[bytes, str]:
+        contents = file.read()
+        if not contents:
+            raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
+        file_path = f"{UPLOAD_DIR}/{filename}"
+        with open(file_path, "wb") as f:
+            f.write(contents)
+        return contents, file_path
 
-    def _initialize_s3(self) -> None:
-        """Initializes the S3 client and bucket name if using S3 storage."""
+    @staticmethod
+    def get_file(file_path: str) -> str:
+        """Handles downloading of the file from local storage."""
+        return file_path
+
+    @staticmethod
+    def delete_file(file_path: str) -> None:
+        """Handles deletion of the file from local storage."""
+        filename = file_path.split("/")[-1]
+        file_path = f"{UPLOAD_DIR}/{filename}"
+        if os.path.isfile(file_path):
+            os.remove(file_path)
+        else:
+            print(f"File {file_path} not found in local storage.")
+
+    @staticmethod
+    def delete_all_files() -> None:
+        """Handles deletion of all files from local storage."""
+        if os.path.exists(UPLOAD_DIR):
+            for filename in os.listdir(UPLOAD_DIR):
+                file_path = os.path.join(UPLOAD_DIR, filename)
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.unlink(file_path)  # Remove the file or link
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)  # Remove the directory
+                except Exception as e:
+                    print(f"Failed to delete {file_path}. Reason: {e}")
+        else:
+            print(f"Directory {UPLOAD_DIR} not found in local storage.")
+
+
+class S3StorageProvider(StorageProvider):
+    def __init__(self):
         self.s3_client = boto3.client(
             "s3",
             region_name=S3_REGION_NAME,
@@ -44,11 +89,9 @@ class StorageProvider:
         )
         self.bucket_name = S3_BUCKET_NAME
 
-    def _upload_to_s3(self, file_path: str, filename: str) -> Tuple[bytes, str]:
+    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
         """Handles uploading of the file to S3 storage."""
-        if not self.s3_client:
-            raise RuntimeError("S3 Client is not initialized.")
-
+        _, file_path = LocalStorageProvider.upload_file(file, filename)
         try:
             self.s3_client.upload_file(file_path, self.bucket_name, filename)
             return (
@@ -58,18 +101,8 @@ class StorageProvider:
         except ClientError as e:
             raise RuntimeError(f"Error uploading file to S3: {e}")
 
-    def _upload_to_local(self, contents: bytes, filename: str) -> Tuple[bytes, str]:
-        """Handles uploading of the file to local storage."""
-        file_path = f"{UPLOAD_DIR}/{filename}"
-        with open(file_path, "wb") as f:
-            f.write(contents)
-        return contents, file_path
-
-    def _get_file_from_s3(self, file_path: str) -> str:
+    def get_file(self, file_path: str) -> str:
         """Handles downloading of the file from S3 storage."""
-        if not self.s3_client:
-            raise RuntimeError("S3 Client is not initialized.")
-
         try:
             bucket_name, key = file_path.split("//")[1].split("/")
             local_file_path = f"{UPLOAD_DIR}/{key}"
@@ -78,33 +111,19 @@ class StorageProvider:
         except ClientError as e:
             raise RuntimeError(f"Error downloading file from S3: {e}")
 
-    def _get_file_from_local(self, file_path: str) -> str:
-        """Handles downloading of the file from local storage."""
-        return file_path
-
-    def _delete_from_s3(self, filename: str) -> None:
+    def delete_file(self, file_path: str) -> None:
         """Handles deletion of the file from S3 storage."""
-        if not self.s3_client:
-            raise RuntimeError("S3 Client is not initialized.")
-
+        filename = file_path.split("/")[-1]
         try:
             self.s3_client.delete_object(Bucket=self.bucket_name, Key=filename)
         except ClientError as e:
             raise RuntimeError(f"Error deleting file from S3: {e}")
 
-    def _delete_from_local(self, filename: str) -> None:
-        """Handles deletion of the file from local storage."""
-        file_path = f"{UPLOAD_DIR}/{filename}"
-        if os.path.isfile(file_path):
-            os.remove(file_path)
-        else:
-            print(f"File {file_path} not found in local storage.")
+        # Always delete from local storage
+        LocalStorageProvider.delete_file(file_path)
 
-    def _delete_all_from_s3(self) -> None:
+    def delete_all_files(self) -> None:
         """Handles deletion of all files from S3 storage."""
-        if not self.s3_client:
-            raise RuntimeError("S3 Client is not initialized.")
-
         try:
             response = self.s3_client.list_objects_v2(Bucket=self.bucket_name)
             if "Contents" in response:
@@ -115,55 +134,13 @@ class StorageProvider:
         except ClientError as e:
             raise RuntimeError(f"Error deleting all files from S3: {e}")
 
-    def _delete_all_from_local(self) -> None:
-        """Handles deletion of all files from local storage."""
-        if os.path.exists(UPLOAD_DIR):
-            for filename in os.listdir(UPLOAD_DIR):
-                file_path = os.path.join(UPLOAD_DIR, filename)
-                try:
-                    if os.path.isfile(file_path) or os.path.islink(file_path):
-                        os.unlink(file_path)  # Remove the file or link
-                    elif os.path.isdir(file_path):
-                        shutil.rmtree(file_path)  # Remove the directory
-                except Exception as e:
-                    print(f"Failed to delete {file_path}. Reason: {e}")
-        else:
-            print(f"Directory {UPLOAD_DIR} not found in local storage.")
-
-    def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
-        """Uploads a file either to S3 or the local file system."""
-        contents = file.read()
-        if not contents:
-            raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
-        contents, file_path = self._upload_to_local(contents, filename)
-
-        if self.storage_provider == "s3":
-            return self._upload_to_s3(file_path, filename)
-        return contents, file_path
-
-    def get_file(self, file_path: str) -> str:
-        """Downloads a file either from S3 or the local file system and returns the file path."""
-        if self.storage_provider == "s3":
-            return self._get_file_from_s3(file_path)
-        return self._get_file_from_local(file_path)
-
-    def delete_file(self, file_path: str) -> None:
-        """Deletes a file either from S3 or the local file system."""
-        filename = file_path.split("/")[-1]
-
-        if self.storage_provider == "s3":
-            self._delete_from_s3(filename)
-
-        # Always delete from local storage
-        self._delete_from_local(filename)
-
-    def delete_all_files(self) -> None:
-        """Deletes all files from the storage."""
-        if self.storage_provider == "s3":
-            self._delete_all_from_s3()
-
         # Always delete from local storage
-        self._delete_all_from_local()
+        LocalStorageProvider.delete_all_files()
 
 
-Storage = StorageProvider(provider=STORAGE_PROVIDER)
+if STORAGE_PROVIDER == "local":
+    Storage = LocalStorageProvider()
+elif STORAGE_PROVIDER == "s3":
+    Storage = S3StorageProvider()
+else:
+    raise RuntimeError(f"Unsupported storage provider: {STORAGE_PROVIDER}")