From 85924e898ec8a7fcf31f457624828acd071b9365 Mon Sep 17 00:00:00 2001
From: The Wind <haifzhu@users.noreply.github.com>
Date: Fri, 28 Feb 2025 15:00:00 +0800
Subject: [PATCH] Fix: enhance aliyun oss access with adding prefix path
 (#5475)

### What problem does this PR solve?

Enhance aliyun oss access with adding prefix path.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 docker/service_conf.yaml.template |  1 +
 rag/utils/oss_conn.py             | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template
index ce41f8147..4dacd8d7f 100644
--- a/docker/service_conf.yaml.template
+++ b/docker/service_conf.yaml.template
@@ -43,6 +43,7 @@ redis:
 #   endpoint_url: '${ENDPOINT}'
 #   region: '${REGION}'
 #   bucket: '${BUCKET}'
+#   prefix_path: '${OSS_PREFIX_PATH}'
 # azure:
 #   auth_type: 'sas'
 #   container_url: 'container_url'
diff --git a/rag/utils/oss_conn.py b/rag/utils/oss_conn.py
index 5525ecbc8..298c1a5b8 100644
--- a/rag/utils/oss_conn.py
+++ b/rag/utils/oss_conn.py
@@ -33,6 +33,7 @@ class RAGFlowOSS(object):
         self.endpoint_url = self.oss_config.get('endpoint_url', None)
         self.region = self.oss_config.get('region', None)
         self.bucket = self.oss_config.get('bucket', None)
+        self.prefix_path = self.oss_config.get('prefix_path', None)
         self.__open__()
 
     @staticmethod
@@ -42,6 +43,14 @@ class RAGFlowOSS(object):
             actual_bucket = self.bucket if self.bucket else bucket
             return method(self, actual_bucket, *args, **kwargs)
         return wrapper
+    
+    @staticmethod
+    def use_prefix_path(method):
+        def wrapper(self, bucket, fnm, *args, **kwargs):
+            # If the prefix path is set, use the prefix path
+            fnm = f"{self.prefix_path}/{fnm}" if self.prefix_path else fnm
+            return method(self, bucket, fnm, *args, **kwargs)
+        return wrapper
 
     def __open__(self):
         try:
@@ -79,8 +88,9 @@ class RAGFlowOSS(object):
         return exists
 
     def health(self):
-        bucket, fnm, binary = "txtxtxtxt1", "txtxtxtxt1", b"_t@@@1"
-
+        bucket = self.bucket
+        fnm = "txtxtxtxt1"
+        fnm, binary = f"{self.prefix_path}/{fnm}" if self.prefix_path else fnm, b"_t@@@1"
         if not self.bucket_exists(bucket):
             self.conn.create_bucket(Bucket=bucket)
             logging.debug(f"create bucket {bucket} ********")
@@ -94,6 +104,7 @@ class RAGFlowOSS(object):
     def list(self, bucket, dir, recursive=True):
         return []
 
+    @use_prefix_path
     @use_default_bucket
     def put(self, bucket, fnm, binary):
         logging.debug(f"bucket name {bucket}; filename :{fnm}:")
@@ -110,6 +121,7 @@ class RAGFlowOSS(object):
                 self.__open__()
                 time.sleep(1)
 
+    @use_prefix_path
     @use_default_bucket
     def rm(self, bucket, fnm):
         try:
@@ -117,6 +129,7 @@ class RAGFlowOSS(object):
         except Exception:
             logging.exception(f"Fail rm {bucket}/{fnm}")
 
+    @use_prefix_path
     @use_default_bucket
     def get(self, bucket, fnm):
         for _ in range(1):
@@ -130,6 +143,7 @@ class RAGFlowOSS(object):
                 time.sleep(1)
         return
 
+    @use_prefix_path
     @use_default_bucket
     def obj_exist(self, bucket, fnm):
         try:
@@ -141,6 +155,7 @@ class RAGFlowOSS(object):
             else:
                 raise
 
+    @use_prefix_path
     @use_default_bucket
     def get_presigned_url(self, bucket, fnm, expires):
         for _ in range(10):
@@ -156,3 +171,4 @@ class RAGFlowOSS(object):
                 self.__open__()
                 time.sleep(1)
         return
+