From 80b9507e7acf2df7ac34f37e89e8a52721ed2958 Mon Sep 17 00:00:00 2001 From: Henrybit Date: Thu, 25 Apr 2024 16:57:19 +0800 Subject: [PATCH] feat: add aliyun oss storage (#3690) Co-authored-by: henrybit --- api/.env.example | 5 +++++ api/config.py | 4 ++++ api/extensions/ext_storage.py | 25 ++++++++++++++++++++++++- api/requirements.txt | 1 + 4 files changed, 34 insertions(+), 1 deletion(-) diff --git a/api/.env.example b/api/.env.example index a228335a22..d8c3f0e129 100644 --- a/api/.env.example +++ b/api/.env.example @@ -52,6 +52,11 @@ AZURE_BLOB_ACCOUNT_NAME=your-account-name AZURE_BLOB_ACCOUNT_KEY=your-account-key AZURE_BLOB_CONTAINER_NAME=yout-container-name AZURE_BLOB_ACCOUNT_URL=https://.blob.core.windows.net +# Aliyun oss Storage configuration +ALIYUN_OSS_BUCKET_NAME=your-bucket-name +ALIYUN_OSS_ACCESS_KEY=your-access-key +ALIYUN_OSS_SECRET_KEY=your-secret-key +ALIYUN_OSS_ENDPOINT=your-endpoint # CORS configuration WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,* diff --git a/api/config.py b/api/config.py index 043afeb579..a269bd582f 100644 --- a/api/config.py +++ b/api/config.py @@ -208,6 +208,10 @@ class Config: self.AZURE_BLOB_ACCOUNT_KEY = get_env('AZURE_BLOB_ACCOUNT_KEY') self.AZURE_BLOB_CONTAINER_NAME = get_env('AZURE_BLOB_CONTAINER_NAME') self.AZURE_BLOB_ACCOUNT_URL = get_env('AZURE_BLOB_ACCOUNT_URL') + self.ALIYUN_OSS_BUCKET_NAME=get_env('ALIYUN_OSS_BUCKET_NAME') + self.ALIYUN_OSS_ACCESS_KEY=get_env('ALIYUN_OSS_ACCESS_KEY') + self.ALIYUN_OSS_SECRET_KEY=get_env('ALIYUN_OSS_SECRET_KEY') + self.ALIYUN_OSS_ENDPOINT=get_env('ALIYUN_OSS_ENDPOINT') # ------------------------ # Vector Store Configurations. diff --git a/api/extensions/ext_storage.py b/api/extensions/ext_storage.py index c2104cdf55..439b485152 100644 --- a/api/extensions/ext_storage.py +++ b/api/extensions/ext_storage.py @@ -6,6 +6,7 @@ from datetime import datetime, timedelta, timezone from typing import Union import boto3 +import oss2 as aliyun_s3 from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas from botocore.client import Config from botocore.exceptions import ClientError @@ -42,7 +43,14 @@ class Storage: ) self.client = BlobServiceClient(account_url=app.config.get('AZURE_BLOB_ACCOUNT_URL'), credential=sas_token) - + elif self.storage_type == 'aliyun-oss': + self.bucket_name = app.config.get('ALIYUN_OSS_BUCKET_NAME') + self.client = aliyun_s3.Bucket( + aliyun_s3.Auth(app.config.get('ALIYUN_OSS_ACCESS_KEY'), app.config.get('ALIYUN_OSS_SECRET_KEY')), + app.config.get('ALIYUN_OSS_ENDPOINT'), + self.bucket_name, + connect_timeout=30 + ) else: self.folder = app.config.get('STORAGE_LOCAL_PATH') if not os.path.isabs(self.folder): @@ -54,6 +62,8 @@ class Storage: elif self.storage_type == 'azure-blob': blob_container = self.client.get_container_client(container=self.bucket_name) blob_container.upload_blob(filename, data) + elif self.storage_type == 'aliyun-oss': + self.client.put_object(filename, data) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -86,6 +96,9 @@ class Storage: blob = self.client.get_container_client(container=self.bucket_name) blob = blob.get_blob_client(blob=filename) data = blob.download_blob().readall() + elif self.storage_type == 'aliyun-oss': + with closing(self.client.get_object(filename)) as obj: + data = obj.read() else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -118,6 +131,10 @@ class Storage: with closing(blob.download_blob()) as blob_stream: while chunk := blob_stream.readall(4096): yield chunk + elif self.storage_type == 'aliyun-oss': + with closing(self.client.get_object(filename)) as obj: + while chunk := obj.read(4096): + yield chunk else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -142,6 +159,8 @@ class Storage: with open(target_filepath, "wb") as my_blob: blob_data = blob.download_blob() blob_data.readinto(my_blob) + elif self.storage_type == 'aliyun-oss': + self.client.get_object_to_file(filename, target_filepath) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -164,6 +183,8 @@ class Storage: elif self.storage_type == 'azure-blob': blob = self.client.get_blob_client(container=self.bucket_name, blob=filename) return blob.exists() + elif self.storage_type == 'aliyun-oss': + return self.client.object_exists(filename) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -178,6 +199,8 @@ class Storage: elif self.storage_type == 'azure-blob': blob_container = self.client.get_container_client(container=self.bucket_name) blob_container.delete_blob(filename) + elif self.storage_type == 'aliyun-oss': + self.client.delete_object(filename) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename diff --git a/api/requirements.txt b/api/requirements.txt index c63b1fcf6d..5984405a41 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -80,3 +80,4 @@ lxml==5.1.0 xlrd~=2.0.1 pydantic~=1.10.0 pgvecto-rs==0.1.4 +oss2==2.15.0 \ No newline at end of file