| import logging |
| from collections.abc import Generator |
| from typing import Union |
|
|
| from flask import Flask |
|
|
| from configs import dify_config |
| from extensions.storage.base_storage import BaseStorage |
| from extensions.storage.storage_type import StorageType |
|
|
|
|
| class Storage: |
| def __init__(self): |
| self.storage_runner = None |
|
|
| def init_app(self, app: Flask): |
| storage_factory = self.get_storage_factory(dify_config.STORAGE_TYPE) |
| with app.app_context(): |
| self.storage_runner = storage_factory() |
|
|
| @staticmethod |
| def get_storage_factory(storage_type: str) -> type[BaseStorage]: |
| match storage_type: |
| case StorageType.S3: |
| from extensions.storage.aws_s3_storage import AwsS3Storage |
|
|
| return AwsS3Storage |
| case StorageType.AZURE_BLOB: |
| from extensions.storage.azure_blob_storage import AzureBlobStorage |
|
|
| return AzureBlobStorage |
| case StorageType.ALIYUN_OSS: |
| from extensions.storage.aliyun_oss_storage import AliyunOssStorage |
|
|
| return AliyunOssStorage |
| case StorageType.GOOGLE_STORAGE: |
| from extensions.storage.google_cloud_storage import GoogleCloudStorage |
|
|
| return GoogleCloudStorage |
| case StorageType.TENCENT_COS: |
| from extensions.storage.tencent_cos_storage import TencentCosStorage |
|
|
| return TencentCosStorage |
| case StorageType.OCI_STORAGE: |
| from extensions.storage.oracle_oci_storage import OracleOCIStorage |
|
|
| return OracleOCIStorage |
| case StorageType.HUAWEI_OBS: |
| from extensions.storage.huawei_obs_storage import HuaweiObsStorage |
|
|
| return HuaweiObsStorage |
| case StorageType.BAIDU_OBS: |
| from extensions.storage.baidu_obs_storage import BaiduObsStorage |
|
|
| return BaiduObsStorage |
| case StorageType.VOLCENGINE_TOS: |
| from extensions.storage.volcengine_tos_storage import VolcengineTosStorage |
|
|
| return VolcengineTosStorage |
| case StorageType.SUPBASE: |
| from extensions.storage.supabase_storage import SupabaseStorage |
|
|
| return SupabaseStorage |
| case StorageType.LOCAL | _: |
| from extensions.storage.local_fs_storage import LocalFsStorage |
|
|
| return LocalFsStorage |
|
|
| def save(self, filename, data): |
| try: |
| self.storage_runner.save(filename, data) |
| except Exception as e: |
| logging.exception("Failed to save file: %s", e) |
| raise e |
|
|
| def load(self, filename: str, /, *, stream: bool = False) -> Union[bytes, Generator]: |
| try: |
| if stream: |
| return self.load_stream(filename) |
| else: |
| return self.load_once(filename) |
| except Exception as e: |
| logging.exception("Failed to load file: %s", e) |
| raise e |
|
|
| def load_once(self, filename: str) -> bytes: |
| try: |
| return self.storage_runner.load_once(filename) |
| except Exception as e: |
| logging.exception("Failed to load_once file: %s", e) |
| raise e |
|
|
| def load_stream(self, filename: str) -> Generator: |
| try: |
| return self.storage_runner.load_stream(filename) |
| except Exception as e: |
| logging.exception("Failed to load_stream file: %s", e) |
| raise e |
|
|
| def download(self, filename, target_filepath): |
| try: |
| self.storage_runner.download(filename, target_filepath) |
| except Exception as e: |
| logging.exception("Failed to download file: %s", e) |
| raise e |
|
|
| def exists(self, filename): |
| try: |
| return self.storage_runner.exists(filename) |
| except Exception as e: |
| logging.exception("Failed to check file exists: %s", e) |
| raise e |
|
|
| def delete(self, filename): |
| try: |
| return self.storage_runner.delete(filename) |
| except Exception as e: |
| logging.exception("Failed to delete file: %s", e) |
| raise e |
|
|
|
|
| storage = Storage() |
|
|
|
|
| def init_app(app: Flask): |
| storage.init_app(app) |
|
|
