| import os |
| import boto3 |
| from botocore.exceptions import ClientError |
| import shutil |
|
|
|
|
| from typing import BinaryIO, Tuple, Optional, Union |
|
|
| from open_webui.constants import ERROR_MESSAGES |
| from open_webui.config import ( |
| STORAGE_PROVIDER, |
| S3_ACCESS_KEY_ID, |
| S3_SECRET_ACCESS_KEY, |
| S3_BUCKET_NAME, |
| S3_REGION_NAME, |
| S3_ENDPOINT_URL, |
| UPLOAD_DIR, |
| ) |
|
|
|
|
| import boto3 |
| from botocore.exceptions import ClientError |
| from typing import BinaryIO, Tuple, Optional |
|
|
|
|
| class StorageProvider: |
| def __init__(self, provider: Optional[str] = None): |
| self.storage_provider: str = provider or STORAGE_PROVIDER |
|
|
| self.s3_client = None |
| self.s3_bucket_name: Optional[str] = None |
|
|
| if self.storage_provider == "s3": |
| self._initialize_s3() |
|
|
| def _initialize_s3(self) -> None: |
| """Initializes the S3 client and bucket name if using S3 storage.""" |
| self.s3_client = boto3.client( |
| "s3", |
| region_name=S3_REGION_NAME, |
| endpoint_url=S3_ENDPOINT_URL, |
| aws_access_key_id=S3_ACCESS_KEY_ID, |
| aws_secret_access_key=S3_SECRET_ACCESS_KEY, |
| ) |
| self.bucket_name = S3_BUCKET_NAME |
|
|
| def _upload_to_s3(self, file_path: str, 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.") |
|
|
| try: |
| self.s3_client.upload_file(file_path, self.bucket_name, filename) |
| return open(file_path, "rb").read(), file_path |
| 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: |
| """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}" |
| self.s3_client.download_file(bucket_name, key, local_file_path) |
| return local_file_path |
| 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: |
| """Handles deletion of the file from S3 storage.""" |
| if not self.s3_client: |
| raise RuntimeError("S3 Client is not initialized.") |
|
|
| 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.") |
|
|
| def _delete_all_from_s3(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: |
| for content in response["Contents"]: |
| self.s3_client.delete_object( |
| Bucket=self.bucket_name, Key=content["Key"] |
| ) |
| 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) |
| elif os.path.isdir(file_path): |
| shutil.rmtree(file_path) |
| 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, filename: str) -> None: |
| """Deletes a file either from S3 or the local file system.""" |
| if self.storage_provider == "s3": |
| self._delete_from_s3(filename) |
|
|
| |
| 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() |
|
|
| |
| self._delete_all_from_local() |
|
|
|
|
| Storage = StorageProvider(provider=STORAGE_PROVIDER) |
|
|
