| from collections.abc import Generator |
|
|
| import oss2 as aliyun_s3 |
|
|
| from configs import dify_config |
| from extensions.storage.base_storage import BaseStorage |
|
|
|
|
| class AliyunOssStorage(BaseStorage): |
| """Implementation for Aliyun OSS storage.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.bucket_name = dify_config.ALIYUN_OSS_BUCKET_NAME |
| self.folder = dify_config.ALIYUN_OSS_PATH |
| oss_auth_method = aliyun_s3.Auth |
| region = None |
| if dify_config.ALIYUN_OSS_AUTH_VERSION == "v4": |
| oss_auth_method = aliyun_s3.AuthV4 |
| region = dify_config.ALIYUN_OSS_REGION |
| oss_auth = oss_auth_method(dify_config.ALIYUN_OSS_ACCESS_KEY, dify_config.ALIYUN_OSS_SECRET_KEY) |
| self.client = aliyun_s3.Bucket( |
| oss_auth, |
| dify_config.ALIYUN_OSS_ENDPOINT, |
| self.bucket_name, |
| connect_timeout=30, |
| region=region, |
| ) |
|
|
| def save(self, filename, data): |
| self.client.put_object(self.__wrapper_folder_filename(filename), data) |
|
|
| def load_once(self, filename: str) -> bytes: |
| obj = self.client.get_object(self.__wrapper_folder_filename(filename)) |
| data = obj.read() |
| return data |
|
|
| def load_stream(self, filename: str) -> Generator: |
| obj = self.client.get_object(self.__wrapper_folder_filename(filename)) |
| while chunk := obj.read(4096): |
| yield chunk |
|
|
| def download(self, filename, target_filepath): |
| self.client.get_object_to_file(self.__wrapper_folder_filename(filename), target_filepath) |
|
|
| def exists(self, filename): |
| return self.client.object_exists(self.__wrapper_folder_filename(filename)) |
|
|
| def delete(self, filename): |
| self.client.delete_object(self.__wrapper_folder_filename(filename)) |
|
|
| def __wrapper_folder_filename(self, filename) -> str: |
| if self.folder: |
| if self.folder.endswith("/"): |
| filename = self.folder + filename |
| else: |
| filename = self.folder + "/" + filename |
| return filename |
|
|
