File size: 4,872 Bytes
20f348c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import logging
from collections.abc import Callable, Generator
from typing import Literal, Union, overload
from flask import Flask
from configs import dify_config
from dify_app import DifyApp
from extensions.storage.base_storage import BaseStorage
from extensions.storage.storage_type import StorageType
logger = logging.getLogger(__name__)
class Storage:
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) -> Callable[[], BaseStorage]:
match storage_type:
case StorageType.S3:
from extensions.storage.aws_s3_storage import AwsS3Storage
return AwsS3Storage
case StorageType.OPENDAL:
from extensions.storage.opendal_storage import OpenDALStorage
return lambda: OpenDALStorage(dify_config.OPENDAL_SCHEME)
case StorageType.LOCAL:
from extensions.storage.opendal_storage import OpenDALStorage
return lambda: OpenDALStorage(scheme="fs", root=dify_config.STORAGE_LOCAL_PATH)
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 _:
raise ValueError(f"unsupported storage type {storage_type}")
def save(self, filename, data):
try:
self.storage_runner.save(filename, data)
except Exception as e:
logger.exception(f"Failed to save file {filename}")
raise e
@overload
def load(self, filename: str, /, *, stream: Literal[False] = False) -> bytes: ...
@overload
def load(self, filename: str, /, *, stream: Literal[True]) -> Generator: ...
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:
logger.exception(f"Failed to load file {filename}")
raise e
def load_once(self, filename: str) -> bytes:
try:
return self.storage_runner.load_once(filename)
except Exception as e:
logger.exception(f"Failed to load_once file {filename}")
raise e
def load_stream(self, filename: str) -> Generator:
try:
return self.storage_runner.load_stream(filename)
except Exception as e:
logger.exception(f"Failed to load_stream file {filename}")
raise e
def download(self, filename, target_filepath):
try:
self.storage_runner.download(filename, target_filepath)
except Exception as e:
logger.exception(f"Failed to download file {filename}")
raise e
def exists(self, filename):
try:
return self.storage_runner.exists(filename)
except Exception as e:
logger.exception(f"Failed to check file exists {filename}")
raise e
def delete(self, filename):
try:
return self.storage_runner.delete(filename)
except Exception as e:
logger.exception(f"Failed to delete file {filename}")
raise e
storage = Storage()
def init_app(app: DifyApp):
storage.init_app(app)
|