Ver Fonte

feat: Add google storage support (#3887)

Co-authored-by: miendinh <miendinh@users.noreply.github.com>
miendinh há 1 ano atrás
pai
commit
9ad489d133
5 ficheiros alterados com 43 adições e 1 exclusões
  1. 3 0
      api/.env.example
  2. 2 0
      api/config.py
  3. 33 0
      api/extensions/ext_storage.py
  4. 1 0
      api/requirements.txt
  5. 4 1
      docker/docker-compose.yaml

+ 3 - 0
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,*

+ 2 - 0
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.

+ 33 - 0
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

+ 1 - 0
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

+ 4 - 1
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://<your_account_name>.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`.