| from collections.abc import Generator | |
| from qcloud_cos import CosConfig, CosS3Client | |
| from configs import dify_config | |
| from extensions.storage.base_storage import BaseStorage | |
| class TencentCosStorage(BaseStorage): | |
| """Implementation for Tencent Cloud COS storage.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.bucket_name = dify_config.TENCENT_COS_BUCKET_NAME | |
| config = CosConfig( | |
| Region=dify_config.TENCENT_COS_REGION, | |
| SecretId=dify_config.TENCENT_COS_SECRET_ID, | |
| SecretKey=dify_config.TENCENT_COS_SECRET_KEY, | |
| Scheme=dify_config.TENCENT_COS_SCHEME, | |
| ) | |
| self.client = CosS3Client(config) | |
| def save(self, filename, data): | |
| self.client.put_object(Bucket=self.bucket_name, Body=data, Key=filename) | |
| def load_once(self, filename: str) -> bytes: | |
| data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].get_raw_stream().read() | |
| return data | |
| def load_stream(self, filename: str) -> Generator: | |
| response = self.client.get_object(Bucket=self.bucket_name, Key=filename) | |
| yield from response["Body"].get_stream(chunk_size=4096) | |
| def download(self, filename, target_filepath): | |
| response = self.client.get_object(Bucket=self.bucket_name, Key=filename) | |
| response["Body"].get_stream_to_file(target_filepath) | |
| def exists(self, filename): | |
| return self.client.object_exists(Bucket=self.bucket_name, Key=filename) | |
| def delete(self, filename): | |
| self.client.delete_object(Bucket=self.bucket_name, Key=filename) | |