refactor: introduce storage factory and speed up api startup by importing storage client on demand (#9086)
This commit is contained in:
Bowen Liang
GitHub
parent
5fcd614186
commit
b360feb4c1
14 changed files with 85 additions and 51 deletions
|
|
@ -4,16 +4,9 @@ from typing import Union
|
|||
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.aliyun_storage import AliyunStorage
|
||||
from extensions.storage.azure_storage import AzureStorage
|
||||
from extensions.storage.baidu_storage import BaiduStorage
|
||||
from extensions.storage.google_storage import GoogleStorage
|
||||
from extensions.storage.huawei_storage import HuaweiStorage
|
||||
from extensions.storage.local_storage import LocalStorage
|
||||
from extensions.storage.oci_storage import OCIStorage
|
||||
from extensions.storage.s3_storage import S3Storage
|
||||
from extensions.storage.tencent_storage import TencentStorage
|
||||
from extensions.storage.volcengine_storage import VolcengineStorage
|
||||
from configs import dify_config
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
from extensions.storage.storage_type import StorageType
|
||||
|
||||
|
||||
class Storage:
|
||||
|
|
@ -21,27 +14,52 @@ class Storage:
|
|||
self.storage_runner = None
|
||||
|
||||
def init_app(self, app: Flask):
|
||||
storage_type = app.config.get("STORAGE_TYPE")
|
||||
if storage_type == "s3":
|
||||
self.storage_runner = S3Storage(app=app)
|
||||
elif storage_type == "azure-blob":
|
||||
self.storage_runner = AzureStorage(app=app)
|
||||
elif storage_type == "aliyun-oss":
|
||||
self.storage_runner = AliyunStorage(app=app)
|
||||
elif storage_type == "google-storage":
|
||||
self.storage_runner = GoogleStorage(app=app)
|
||||
elif storage_type == "tencent-cos":
|
||||
self.storage_runner = TencentStorage(app=app)
|
||||
elif storage_type == "oci-storage":
|
||||
self.storage_runner = OCIStorage(app=app)
|
||||
elif storage_type == "huawei-obs":
|
||||
self.storage_runner = HuaweiStorage(app=app)
|
||||
elif storage_type == "baidu-obs":
|
||||
self.storage_runner = BaiduStorage(app=app)
|
||||
elif storage_type == "volcengine-tos":
|
||||
self.storage_runner = VolcengineStorage(app=app)
|
||||
else:
|
||||
self.storage_runner = LocalStorage(app=app)
|
||||
storage_factory = self.get_storage_factory(dify_config.STORAGE_TYPE)
|
||||
self.storage_runner = storage_factory(app=app)
|
||||
|
||||
@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.LOCAL | _:
|
||||
from extensions.storage.local_fs_storage import LocalFsStorage
|
||||
|
||||
return LocalFsStorage
|
||||
|
||||
def save(self, filename, data):
|
||||
try:
|
||||
|
|
|
|||
64
api/extensions/storage/aliyun_oss_storage.py
Normal file
64
api/extensions/storage/aliyun_oss_storage.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
from collections.abc import Generator
|
||||
from contextlib import closing
|
||||
|
||||
import oss2 as aliyun_s3
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class AliyunOssStorage(BaseStorage):
|
||||
"""Implementation for Aliyun OSS storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("ALIYUN_OSS_BUCKET_NAME")
|
||||
self.folder = app.config.get("ALIYUN_OSS_PATH")
|
||||
oss_auth_method = aliyun_s3.Auth
|
||||
region = None
|
||||
if app_config.get("ALIYUN_OSS_AUTH_VERSION") == "v4":
|
||||
oss_auth_method = aliyun_s3.AuthV4
|
||||
region = app_config.get("ALIYUN_OSS_REGION")
|
||||
oss_auth = oss_auth_method(app_config.get("ALIYUN_OSS_ACCESS_KEY"), app_config.get("ALIYUN_OSS_SECRET_KEY"))
|
||||
self.client = aliyun_s3.Bucket(
|
||||
oss_auth,
|
||||
app_config.get("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:
|
||||
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj:
|
||||
data = obj.read()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj:
|
||||
while chunk := obj.read(4096):
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
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
|
||||
86
api/extensions/storage/aws_s3_storage.py
Normal file
86
api/extensions/storage/aws_s3_storage.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from collections.abc import Generator
|
||||
from contextlib import closing
|
||||
|
||||
import boto3
|
||||
from botocore.client import Config
|
||||
from botocore.exceptions import ClientError
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class AwsS3Storage(BaseStorage):
|
||||
"""Implementation for Amazon Web Services S3 storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("S3_BUCKET_NAME")
|
||||
if app_config.get("S3_USE_AWS_MANAGED_IAM"):
|
||||
session = boto3.Session()
|
||||
self.client = session.client("s3")
|
||||
else:
|
||||
self.client = boto3.client(
|
||||
"s3",
|
||||
aws_secret_access_key=app_config.get("S3_SECRET_KEY"),
|
||||
aws_access_key_id=app_config.get("S3_ACCESS_KEY"),
|
||||
endpoint_url=app_config.get("S3_ENDPOINT"),
|
||||
region_name=app_config.get("S3_REGION"),
|
||||
config=Config(s3={"addressing_style": app_config.get("S3_ADDRESS_STYLE")}),
|
||||
)
|
||||
# create bucket
|
||||
try:
|
||||
self.client.head_bucket(Bucket=self.bucket_name)
|
||||
except ClientError as e:
|
||||
# if bucket not exists, create it
|
||||
if e.response["Error"]["Code"] == "404":
|
||||
self.client.create_bucket(Bucket=self.bucket_name)
|
||||
# if bucket is not accessible, pass, maybe the bucket is existing but not accessible
|
||||
elif e.response["Error"]["Code"] == "403":
|
||||
pass
|
||||
else:
|
||||
# other error, raise exception
|
||||
raise
|
||||
|
||||
def save(self, filename, data):
|
||||
self.client.put_object(Bucket=self.bucket_name, Key=filename, Body=data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
try:
|
||||
with closing(self.client) as client:
|
||||
data = client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
|
||||
except ClientError as ex:
|
||||
if ex.response["Error"]["Code"] == "NoSuchKey":
|
||||
raise FileNotFoundError("File not found")
|
||||
else:
|
||||
raise
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
try:
|
||||
with closing(self.client) as client:
|
||||
response = client.get_object(Bucket=self.bucket_name, Key=filename)
|
||||
yield from response["Body"].iter_chunks()
|
||||
except ClientError as ex:
|
||||
if ex.response["Error"]["Code"] == "NoSuchKey":
|
||||
raise FileNotFoundError("File not found")
|
||||
else:
|
||||
raise
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
with closing(self.client) as client:
|
||||
client.download_file(self.bucket_name, filename, target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
with closing(self.client) as client:
|
||||
try:
|
||||
client.head_object(Bucket=self.bucket_name, Key=filename)
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete(self, filename):
|
||||
self.client.delete_object(Bucket=self.bucket_name, Key=filename)
|
||||
78
api/extensions/storage/azure_blob_storage.py
Normal file
78
api/extensions/storage/azure_blob_storage.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
from collections.abc import Generator
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas
|
||||
from flask import Flask
|
||||
|
||||
from extensions.ext_redis import redis_client
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class AzureBlobStorage(BaseStorage):
|
||||
"""Implementation for Azure Blob storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("AZURE_BLOB_CONTAINER_NAME")
|
||||
self.account_url = app_config.get("AZURE_BLOB_ACCOUNT_URL")
|
||||
self.account_name = app_config.get("AZURE_BLOB_ACCOUNT_NAME")
|
||||
self.account_key = app_config.get("AZURE_BLOB_ACCOUNT_KEY")
|
||||
|
||||
def save(self, filename, data):
|
||||
client = self._sync_client()
|
||||
blob_container = client.get_container_client(container=self.bucket_name)
|
||||
blob_container.upload_blob(filename, data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
client = self._sync_client()
|
||||
blob = client.get_container_client(container=self.bucket_name)
|
||||
blob = blob.get_blob_client(blob=filename)
|
||||
data = blob.download_blob().readall()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
client = self._sync_client()
|
||||
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
blob = client.get_blob_client(container=self.bucket_name, blob=filename)
|
||||
blob_data = blob.download_blob()
|
||||
yield from blob_data.chunks()
|
||||
|
||||
return generate(filename)
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
client = self._sync_client()
|
||||
|
||||
blob = client.get_blob_client(container=self.bucket_name, blob=filename)
|
||||
with open(target_filepath, "wb") as my_blob:
|
||||
blob_data = blob.download_blob()
|
||||
blob_data.readinto(my_blob)
|
||||
|
||||
def exists(self, filename):
|
||||
client = self._sync_client()
|
||||
|
||||
blob = client.get_blob_client(container=self.bucket_name, blob=filename)
|
||||
return blob.exists()
|
||||
|
||||
def delete(self, filename):
|
||||
client = self._sync_client()
|
||||
|
||||
blob_container = client.get_container_client(container=self.bucket_name)
|
||||
blob_container.delete_blob(filename)
|
||||
|
||||
def _sync_client(self):
|
||||
cache_key = "azure_blob_sas_token_{}_{}".format(self.account_name, self.account_key)
|
||||
cache_result = redis_client.get(cache_key)
|
||||
if cache_result is not None:
|
||||
sas_token = cache_result.decode("utf-8")
|
||||
else:
|
||||
sas_token = generate_account_sas(
|
||||
account_name=self.account_name,
|
||||
account_key=self.account_key,
|
||||
resource_types=ResourceTypes(service=True, container=True, object=True),
|
||||
permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True),
|
||||
expiry=datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(hours=1),
|
||||
)
|
||||
redis_client.set(cache_key, sas_token, ex=3000)
|
||||
return BlobServiceClient(account_url=self.account_url, credential=sas_token)
|
||||
60
api/extensions/storage/baidu_obs_storage.py
Normal file
60
api/extensions/storage/baidu_obs_storage.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
import base64
|
||||
import hashlib
|
||||
from collections.abc import Generator
|
||||
|
||||
from baidubce.auth.bce_credentials import BceCredentials
|
||||
from baidubce.bce_client_configuration import BceClientConfiguration
|
||||
from baidubce.services.bos.bos_client import BosClient
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class BaiduObsStorage(BaseStorage):
|
||||
"""Implementation for Baidu OBS storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("BAIDU_OBS_BUCKET_NAME")
|
||||
client_config = BceClientConfiguration(
|
||||
credentials=BceCredentials(
|
||||
access_key_id=app_config.get("BAIDU_OBS_ACCESS_KEY"),
|
||||
secret_access_key=app_config.get("BAIDU_OBS_SECRET_KEY"),
|
||||
),
|
||||
endpoint=app_config.get("BAIDU_OBS_ENDPOINT"),
|
||||
)
|
||||
|
||||
self.client = BosClient(config=client_config)
|
||||
|
||||
def save(self, filename, data):
|
||||
md5 = hashlib.md5()
|
||||
md5.update(data)
|
||||
content_md5 = base64.standard_b64encode(md5.digest())
|
||||
self.client.put_object(
|
||||
bucket_name=self.bucket_name, key=filename, data=data, content_length=len(data), content_md5=content_md5
|
||||
)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
response = self.client.get_object(bucket_name=self.bucket_name, key=filename)
|
||||
return response.data.read()
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
response = self.client.get_object(bucket_name=self.bucket_name, key=filename).data
|
||||
while chunk := response.read(4096):
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
self.client.get_object_to_file(bucket_name=self.bucket_name, key=filename, file_name=target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
res = self.client.get_object_meta_data(bucket_name=self.bucket_name, key=filename)
|
||||
if res is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def delete(self, filename):
|
||||
self.client.delete_object(bucket_name=self.bucket_name, key=filename)
|
||||
64
api/extensions/storage/google_cloud_storage.py
Normal file
64
api/extensions/storage/google_cloud_storage.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
import base64
|
||||
import io
|
||||
import json
|
||||
from collections.abc import Generator
|
||||
from contextlib import closing
|
||||
|
||||
from flask import Flask
|
||||
from google.cloud import storage as google_cloud_storage
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class GoogleCloudStorage(BaseStorage):
|
||||
"""Implementation for Google Cloud storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("GOOGLE_STORAGE_BUCKET_NAME")
|
||||
service_account_json_str = app_config.get("GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64")
|
||||
# if service_account_json_str is empty, use Application Default Credentials
|
||||
if service_account_json_str:
|
||||
service_account_json = base64.b64decode(service_account_json_str).decode("utf-8")
|
||||
# convert str to object
|
||||
service_account_obj = json.loads(service_account_json)
|
||||
self.client = google_cloud_storage.Client.from_service_account_info(service_account_obj)
|
||||
else:
|
||||
self.client = google_cloud_storage.Client()
|
||||
|
||||
def save(self, filename, data):
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
blob = bucket.blob(filename)
|
||||
with io.BytesIO(data) as stream:
|
||||
blob.upload_from_file(stream)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
blob = bucket.get_blob(filename)
|
||||
data = blob.download_as_bytes()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
blob = bucket.get_blob(filename)
|
||||
with closing(blob.open(mode="rb")) as blob_stream:
|
||||
while chunk := blob_stream.read(4096):
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
blob = bucket.get_blob(filename)
|
||||
blob.download_to_filename(target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
blob = bucket.blob(filename)
|
||||
return blob.exists()
|
||||
|
||||
def delete(self, filename):
|
||||
bucket = self.client.get_bucket(self.bucket_name)
|
||||
bucket.delete_blob(filename)
|
||||
54
api/extensions/storage/huawei_obs_storage.py
Normal file
54
api/extensions/storage/huawei_obs_storage.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
from collections.abc import Generator
|
||||
|
||||
from flask import Flask
|
||||
from obs import ObsClient
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class HuaweiObsStorage(BaseStorage):
|
||||
"""Implementation for Huawei OBS storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("HUAWEI_OBS_BUCKET_NAME")
|
||||
self.client = ObsClient(
|
||||
access_key_id=app_config.get("HUAWEI_OBS_ACCESS_KEY"),
|
||||
secret_access_key=app_config.get("HUAWEI_OBS_SECRET_KEY"),
|
||||
server=app_config.get("HUAWEI_OBS_SERVER"),
|
||||
)
|
||||
|
||||
def save(self, filename, data):
|
||||
self.client.putObject(bucketName=self.bucket_name, objectKey=filename, content=data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
data = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response.read()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
response = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response
|
||||
while chunk := response.read(4096):
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
self.client.getObject(bucketName=self.bucket_name, objectKey=filename, downloadPath=target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
res = self._get_meta(filename)
|
||||
if res is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
def delete(self, filename):
|
||||
self.client.deleteObject(bucketName=self.bucket_name, objectKey=filename)
|
||||
|
||||
def _get_meta(self, filename):
|
||||
res = self.client.getObjectMetadata(bucketName=self.bucket_name, objectKey=filename)
|
||||
if res.status < 300:
|
||||
return res
|
||||
else:
|
||||
return None
|
||||
85
api/extensions/storage/local_fs_storage.py
Normal file
85
api/extensions/storage/local_fs_storage.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
import os
|
||||
import shutil
|
||||
from collections.abc import Generator
|
||||
from pathlib import Path
|
||||
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class LocalFsStorage(BaseStorage):
|
||||
"""Implementation for local filesystem storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
folder = self.app.config.get("STORAGE_LOCAL_PATH")
|
||||
if not os.path.isabs(folder):
|
||||
folder = os.path.join(app.root_path, folder)
|
||||
self.folder = folder
|
||||
|
||||
def save(self, filename, data):
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
|
||||
folder = os.path.dirname(filename)
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
|
||||
Path(os.path.join(os.getcwd(), filename)).write_bytes(data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
|
||||
if not os.path.exists(filename):
|
||||
raise FileNotFoundError("File not found")
|
||||
|
||||
data = Path(filename).read_bytes()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
|
||||
if not os.path.exists(filename):
|
||||
raise FileNotFoundError("File not found")
|
||||
|
||||
with open(filename, "rb") as f:
|
||||
while chunk := f.read(4096): # Read in chunks of 4KB
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
|
||||
if not os.path.exists(filename):
|
||||
raise FileNotFoundError("File not found")
|
||||
|
||||
shutil.copyfile(filename, target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
|
||||
return os.path.exists(filename)
|
||||
|
||||
def delete(self, filename):
|
||||
if not self.folder or self.folder.endswith("/"):
|
||||
filename = self.folder + filename
|
||||
else:
|
||||
filename = self.folder + "/" + filename
|
||||
if os.path.exists(filename):
|
||||
os.remove(filename)
|
||||
67
api/extensions/storage/oracle_oci_storage.py
Normal file
67
api/extensions/storage/oracle_oci_storage.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
from collections.abc import Generator
|
||||
from contextlib import closing
|
||||
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class OracleOCIStorage(BaseStorage):
|
||||
"""Implementation for Oracle OCI storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("OCI_BUCKET_NAME")
|
||||
self.client = boto3.client(
|
||||
"s3",
|
||||
aws_secret_access_key=app_config.get("OCI_SECRET_KEY"),
|
||||
aws_access_key_id=app_config.get("OCI_ACCESS_KEY"),
|
||||
endpoint_url=app_config.get("OCI_ENDPOINT"),
|
||||
region_name=app_config.get("OCI_REGION"),
|
||||
)
|
||||
|
||||
def save(self, filename, data):
|
||||
self.client.put_object(Bucket=self.bucket_name, Key=filename, Body=data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
try:
|
||||
with closing(self.client) as client:
|
||||
data = client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
|
||||
except ClientError as ex:
|
||||
if ex.response["Error"]["Code"] == "NoSuchKey":
|
||||
raise FileNotFoundError("File not found")
|
||||
else:
|
||||
raise
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
try:
|
||||
with closing(self.client) as client:
|
||||
response = client.get_object(Bucket=self.bucket_name, Key=filename)
|
||||
yield from response["Body"].iter_chunks()
|
||||
except ClientError as ex:
|
||||
if ex.response["Error"]["Code"] == "NoSuchKey":
|
||||
raise FileNotFoundError("File not found")
|
||||
else:
|
||||
raise
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
with closing(self.client) as client:
|
||||
client.download_file(self.bucket_name, filename, target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
with closing(self.client) as client:
|
||||
try:
|
||||
client.head_object(Bucket=self.bucket_name, Key=filename)
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def delete(self, filename):
|
||||
self.client.delete_object(Bucket=self.bucket_name, Key=filename)
|
||||
14
api/extensions/storage/storage_type.py
Normal file
14
api/extensions/storage/storage_type.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
from enum import Enum
|
||||
|
||||
|
||||
class StorageType(str, Enum):
|
||||
ALIYUN_OSS = "aliyun-oss"
|
||||
AZURE_BLOB = "azure-blob"
|
||||
BAIDU_OBS = "baidu-obs"
|
||||
GOOGLE_STORAGE = "google-storage"
|
||||
HUAWEI_OBS = "huawei-obs"
|
||||
LOCAL = "local"
|
||||
OCI_STORAGE = "oci-storage"
|
||||
S3 = "s3"
|
||||
TENCENT_COS = "tencent-cos"
|
||||
VOLCENGINE_TOS = "volcengine-tos"
|
||||
46
api/extensions/storage/tencent_cos_storage.py
Normal file
46
api/extensions/storage/tencent_cos_storage.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from collections.abc import Generator
|
||||
|
||||
from flask import Flask
|
||||
from qcloud_cos import CosConfig, CosS3Client
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class TencentCosStorage(BaseStorage):
|
||||
"""Implementation for Tencent Cloud COS storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("TENCENT_COS_BUCKET_NAME")
|
||||
config = CosConfig(
|
||||
Region=app_config.get("TENCENT_COS_REGION"),
|
||||
SecretId=app_config.get("TENCENT_COS_SECRET_ID"),
|
||||
SecretKey=app_config.get("TENCENT_COS_SECRET_KEY"),
|
||||
Scheme=app_config.get("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:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
|
||||
yield from response["Body"].get_stream(chunk_size=4096)
|
||||
|
||||
return generate()
|
||||
|
||||
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)
|
||||
48
api/extensions/storage/volcengine_tos_storage.py
Normal file
48
api/extensions/storage/volcengine_tos_storage.py
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
from collections.abc import Generator
|
||||
|
||||
import tos
|
||||
from flask import Flask
|
||||
|
||||
from extensions.storage.base_storage import BaseStorage
|
||||
|
||||
|
||||
class VolcengineTosStorage(BaseStorage):
|
||||
"""Implementation for Volcengine TOS storage."""
|
||||
|
||||
def __init__(self, app: Flask):
|
||||
super().__init__(app)
|
||||
app_config = self.app.config
|
||||
self.bucket_name = app_config.get("VOLCENGINE_TOS_BUCKET_NAME")
|
||||
self.client = tos.TosClientV2(
|
||||
ak=app_config.get("VOLCENGINE_TOS_ACCESS_KEY"),
|
||||
sk=app_config.get("VOLCENGINE_TOS_SECRET_KEY"),
|
||||
endpoint=app_config.get("VOLCENGINE_TOS_ENDPOINT"),
|
||||
region=app_config.get("VOLCENGINE_TOS_REGION"),
|
||||
)
|
||||
|
||||
def save(self, filename, data):
|
||||
self.client.put_object(bucket=self.bucket_name, key=filename, content=data)
|
||||
|
||||
def load_once(self, filename: str) -> bytes:
|
||||
data = self.client.get_object(bucket=self.bucket_name, key=filename).read()
|
||||
return data
|
||||
|
||||
def load_stream(self, filename: str) -> Generator:
|
||||
def generate(filename: str = filename) -> Generator:
|
||||
response = self.client.get_object(bucket=self.bucket_name, key=filename)
|
||||
while chunk := response.read(4096):
|
||||
yield chunk
|
||||
|
||||
return generate()
|
||||
|
||||
def download(self, filename, target_filepath):
|
||||
self.client.get_object_to_file(bucket=self.bucket_name, key=filename, file_path=target_filepath)
|
||||
|
||||
def exists(self, filename):
|
||||
res = self.client.head_object(bucket=self.bucket_name, key=filename)
|
||||
if res.status_code != 200:
|
||||
return False
|
||||
return True
|
||||
|
||||
def delete(self, filename):
|
||||
self.client.delete_object(bucket=self.bucket_name, key=filename)
|
||||
2
api/poetry.lock
generated
2
api/poetry.lock
generated
|
|
@ -10595,4 +10595,4 @@ cffi = ["cffi (>=1.11)"]
|
|||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = ">=3.10,<3.13"
|
||||
content-hash = "fd183812f910faf4e840267501c571db5d758ad6eb328d106ba6f79a0322a555"
|
||||
content-hash = "34ba8efcc67da342036ef075b693f59fdc67d246f40b857c9c1bd6f80c7283bd"
|
||||
|
|
|
|||
|
|
@ -113,7 +113,6 @@ authlib = "1.3.1"
|
|||
azure-ai-inference = "~1.0.0b3"
|
||||
azure-ai-ml = "~1.20.0"
|
||||
azure-identity = "1.16.1"
|
||||
azure-storage-blob = "12.13.0"
|
||||
beautifulsoup4 = "4.12.2"
|
||||
boto3 = "1.35.17"
|
||||
bs4 = "~0.0.1"
|
||||
|
|
@ -221,6 +220,7 @@ yfinance = "~0.2.40"
|
|||
# Required for storage clients
|
||||
############################################################
|
||||
[tool.poetry.group.storage.dependencies]
|
||||
azure-storage-blob = "12.13.0"
|
||||
bce-python-sdk = "~0.9.23"
|
||||
cos-python-sdk-v5 = "1.9.30"
|
||||
esdk-obs-python = "3.24.6.1"
|
||||
|
|
|
|||
Loading…
Reference in a new issue