Underground-Digital/Workflow-Engine

1import logging2from collections.abc import Generator3from typing import Union4 5from flask import Flask6 7from configs import dify_config8from extensions.storage.base_storage import BaseStorage9from extensions.storage.storage_type import StorageType10 11 12class Storage:13    def __init__(self):14        self.storage_runner = None15 16    def init_app(self, app: Flask):17        storage_factory = self.get_storage_factory(dify_config.STORAGE_TYPE)18        with app.app_context():19            self.storage_runner = storage_factory()20 21    @staticmethod22    def get_storage_factory(storage_type: str) -> type[BaseStorage]:23        match storage_type:24            case StorageType.S3:25                from extensions.storage.aws_s3_storage import AwsS3Storage26 27                return AwsS3Storage28            case StorageType.AZURE_BLOB:29                from extensions.storage.azure_blob_storage import AzureBlobStorage30 31                return AzureBlobStorage32            case StorageType.ALIYUN_OSS:33                from extensions.storage.aliyun_oss_storage import AliyunOssStorage34 35                return AliyunOssStorage36            case StorageType.GOOGLE_STORAGE:37                from extensions.storage.google_cloud_storage import GoogleCloudStorage38 39                return GoogleCloudStorage40            case StorageType.TENCENT_COS:41                from extensions.storage.tencent_cos_storage import TencentCosStorage42 43                return TencentCosStorage44            case StorageType.OCI_STORAGE:45                from extensions.storage.oracle_oci_storage import OracleOCIStorage46 47                return OracleOCIStorage48            case StorageType.HUAWEI_OBS:49                from extensions.storage.huawei_obs_storage import HuaweiObsStorage50 51                return HuaweiObsStorage52            case StorageType.BAIDU_OBS:53                from extensions.storage.baidu_obs_storage import BaiduObsStorage54 55                return BaiduObsStorage56            case StorageType.VOLCENGINE_TOS:57                from extensions.storage.volcengine_tos_storage import VolcengineTosStorage58 59                return VolcengineTosStorage60            case StorageType.SUPBASE:61                from extensions.storage.supabase_storage import SupabaseStorage62 63                return SupabaseStorage64            case StorageType.LOCAL | _:65                from extensions.storage.local_fs_storage import LocalFsStorage66 67                return LocalFsStorage68 69    def save(self, filename, data):70        try:71            self.storage_runner.save(filename, data)72        except Exception as e:73            logging.exception("Failed to save file: %s", e)74            raise e75 76    def load(self, filename: str, /, *, stream: bool = False) -> Union[bytes, Generator]:77        try:78            if stream:79                return self.load_stream(filename)80            else:81                return self.load_once(filename)82        except Exception as e:83            logging.exception("Failed to load file: %s", e)84            raise e85 86    def load_once(self, filename: str) -> bytes:87        try:88            return self.storage_runner.load_once(filename)89        except Exception as e:90            logging.exception("Failed to load_once file: %s", e)91            raise e92 93    def load_stream(self, filename: str) -> Generator:94        try:95            return self.storage_runner.load_stream(filename)96        except Exception as e:97            logging.exception("Failed to load_stream file: %s", e)98            raise e99 100    def download(self, filename, target_filepath):101        try:102            self.storage_runner.download(filename, target_filepath)103        except Exception as e:104            logging.exception("Failed to download file: %s", e)105            raise e106 107    def exists(self, filename):108        try:109            return self.storage_runner.exists(filename)110        except Exception as e:111            logging.exception("Failed to check file exists: %s", e)112            raise e113 114    def delete(self, filename):115        try:116            return self.storage_runner.delete(filename)117        except Exception as e:118            logging.exception("Failed to delete file: %s", e)119            raise e120 121 122storage = Storage()123 124 125def init_app(app: Flask):126    storage.init_app(app)127