From 9ad489d1333545f723fbdfd24be6987f0e7a54b1 Mon Sep 17 00:00:00 2001 From: miendinh <22139872+miendinh@users.noreply.github.com> Date: Sat, 27 Apr 2024 17:26:52 +0700 Subject: [PATCH] feat: Add google storage support (#3887) Co-authored-by: miendinh --- api/.env.example | 3 +++ api/config.py | 2 ++ api/extensions/ext_storage.py | 33 +++++++++++++++++++++++++++++++++ api/requirements.txt | 1 + docker/docker-compose.yaml | 5 ++++- 5 files changed, 43 insertions(+), 1 deletion(-) diff --git a/api/.env.example b/api/.env.example index c61cb60d3e..0c30c341a6 100644 --- a/api/.env.example +++ b/api/.env.example @@ -54,6 +54,9 @@ 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 +# Google Storage configuration +GOOGLE_STORAGE_BUCKET_NAME=yout-bucket-name +GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON=your-google-service-account-json-base64-string # CORS configuration WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,* diff --git a/api/config.py b/api/config.py index 8fb6b83cb3..0de1a01881 100644 --- a/api/config.py +++ b/api/config.py @@ -213,6 +213,8 @@ class Config: 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') + self.GOOGLE_STORAGE_BUCKET_NAME = get_env('GOOGLE_STORAGE_BUCKET_NAME') + self.GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64 = get_env('GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64') # ------------------------ # Vector Store Configurations. diff --git a/api/extensions/ext_storage.py b/api/extensions/ext_storage.py index 439b485152..db66df88d8 100644 --- a/api/extensions/ext_storage.py +++ b/api/extensions/ext_storage.py @@ -1,3 +1,4 @@ +import base64 import os import shutil from collections.abc import Generator @@ -11,6 +12,7 @@ from azure.storage.blob import AccountSasPermissions, BlobServiceClient, Resourc from botocore.client import Config from botocore.exceptions import ClientError from flask import Flask +from google.cloud import storage as GoogleStorage class Storage: @@ -51,6 +53,10 @@ class Storage: self.bucket_name, connect_timeout=30 ) + elif self.storage_type == 'google-storage': + self.bucket_name = app.config.get('GOOGLE_STORAGE_BUCKET_NAME') + service_account_json = base64.b64decode(app.config.get('GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64')).decode('utf-8') + self.client = GoogleStorage.Client().from_service_account_json(service_account_json) else: self.folder = app.config.get('STORAGE_LOCAL_PATH') if not os.path.isabs(self.folder): @@ -64,6 +70,10 @@ class Storage: blob_container.upload_blob(filename, data) elif self.storage_type == 'aliyun-oss': self.client.put_object(filename, data) + elif self.storage_type == 'google-storage': + bucket = self.client.get_bucket(self.bucket_name) + blob = bucket.blob(filename) + blob.upload_from_file(data) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -99,6 +109,10 @@ class Storage: elif self.storage_type == 'aliyun-oss': with closing(self.client.get_object(filename)) as obj: data = obj.read() + elif self.storage_type == 'google-storage': + bucket = self.client.get_bucket(self.bucket_name) + blob = bucket.get_blob(filename) + data = blob.download_as_bytes() else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -135,6 +149,12 @@ class Storage: with closing(self.client.get_object(filename)) as obj: while chunk := obj.read(4096): yield chunk + elif self.storage_type == 'google-storage': + 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 else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -161,6 +181,12 @@ class Storage: blob_data.readinto(my_blob) elif self.storage_type == 'aliyun-oss': self.client.get_object_to_file(filename, target_filepath) + elif self.storage_type == 'google-storage': + bucket = self.client.get_bucket(self.bucket_name) + blob = bucket.get_blob(filename) + with open(target_filepath, "wb") as my_blob: + blob_data = blob.download_blob() + blob_data.readinto(my_blob) else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -185,6 +211,10 @@ class Storage: return blob.exists() elif self.storage_type == 'aliyun-oss': return self.client.object_exists(filename) + elif self.storage_type == 'google-storage': + bucket = self.client.get_bucket(self.bucket_name) + blob = bucket.blob(filename) + return blob.exists() else: if not self.folder or self.folder.endswith('/'): filename = self.folder + filename @@ -201,6 +231,9 @@ class Storage: blob_container.delete_blob(filename) elif self.storage_type == 'aliyun-oss': self.client.delete_object(filename) + elif self.storage_type == 'google-storage': + bucket = self.client.get_bucket(self.bucket_name) + bucket.delete_blob(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 5984405a41..bcf248a4c4 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -44,6 +44,7 @@ google-auth-httplib2==0.2.0 google-generativeai==0.5.0 google-search-results==2.4.2 googleapis-common-protos==1.63.0 +google-cloud-storage==2.16.0 replicate~=0.22.0 websocket-client~=1.7.0 dashscope[tokenizer]~=1.17.0 diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 46ed637424..056dc10630 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -70,7 +70,7 @@ services: # If you want to enable cross-origin support, # you must use the HTTPS protocol and set the configuration to `SameSite=None, Secure=true, HttpOnly=true`. # - # The type of storage to use for storing user files. Supported values are `local` and `s3` and `azure-blob`, Default: `local` + # The type of storage to use for storing user files. Supported values are `local` and `s3` and `azure-blob` and `google-storage`, Default: `local` STORAGE_TYPE: local # The path to the local storage directory, the directory relative the root path of API service codes or absolute path. Default: `storage` or `/home/john/storage`. # only available when STORAGE_TYPE is `local`. @@ -86,6 +86,9 @@ services: AZURE_BLOB_ACCOUNT_KEY: 'difyai' AZURE_BLOB_CONTAINER_NAME: 'difyai-container' AZURE_BLOB_ACCOUNT_URL: 'https://.blob.core.windows.net' + # The Google storage configurations, only available when STORAGE_TYPE is `google-storage`. + GOOGLE_STORAGE_BUCKET_NAME: 'yout-bucket-name' + GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64: 'your-google-service-account-json-base64-string' # The type of vector store to use. Supported values are `weaviate`, `qdrant`, `milvus`, `relyt`. VECTOR_STORE: weaviate # The Weaviate endpoint URL. Only available when VECTOR_STORE is `weaviate`.