Dev: Clean code.

infiniflow · Sep 5, 2024 · faea77c · faea77c
1 parent 411cb47
commit faea77c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 18 deletions.
diff --git a/rag/utils/azure_sas_conn.py b/rag/utils/azure_sas_conn.py
@@ -1,7 +1,6 @@
 import os
 import time
 from io import BytesIO
-from rag import settings
 from rag.settings import azure_logger
 from rag.utils import singleton
 from azure.storage.blob import ContainerClient
@@ -11,6 +10,8 @@
 class RAGFlowAzureSasBlob(object):
     def __init__(self):
         self.conn = None
+        self.account_url = os.getenv('CONTAINER_URL', None)
+        self.sas_token = os.getenv('SAS_TOKEN', None)
         self.__open__()
 
     def __open__(self):
@@ -21,12 +22,10 @@ def __open__(self):
             pass
 
         try:
-            account_url = os.getenv('CONTAINER_URL', None)
-            sas_token = os.getenv('SAS_TOKEN', None)
-            self.conn = ContainerClient.from_container_url(account_url + "?" + sas_token)
+            self.conn = ContainerClient.from_container_url(self.account_url + "?" + self.sas_token)
         except Exception as e:
             azure_logger.error(
-                "Fail to connect %s " % settings.MINIO["host"] + str(e))
+                "Fail to connect %s " % self.account_url + str(e))
 
     def __close__(self):
         del self.conn

diff --git a/rag/utils/azure_spn_conn.py b/rag/utils/azure_spn_conn.py
@@ -1,6 +1,5 @@
 import os
 import time
-from rag import settings
 from rag.settings import azure_logger
 from rag.utils import singleton
 from azure.identity import ClientSecretCredential, AzureAuthorityHosts
@@ -30,7 +29,7 @@ def __open__(self):
             self.conn = FileSystemClient(account_url=self.account_url, file_system_name=self.container_name, credential=credentials)
         except Exception as e:
             azure_logger.error(
-                "Fail to connect %s " % settings.MINIO["host"] + str(e))
+                "Fail to connect %s " % self.account_url + str(e))
 
     def __close__(self):
         del self.conn

diff --git a/rag/utils/s3_conn.py b/rag/utils/s3_conn.py
@@ -1,17 +1,20 @@
 import boto3
+import os
 from botocore.exceptions import ClientError
 from botocore.client import Config
 import time
 from io import BytesIO
-from rag import settings
 from rag.settings import s3_logger
 from rag.utils import singleton
 
-
 @singleton
 class RAGFlowS3(object):
     def __init__(self):
         self.conn = None
+        self.endpoint = os.getenv('ENDPOINT', None)
+        self.access_key = os.getenv('ACCESS_KEY', None)
+        self.secret_key = os.getenv('SECRET_KEY', None)
+        self.region = os.getenv('REGION', None)
         self.__open__()
 
     def __open__(self):
@@ -22,10 +25,6 @@ def __open__(self):
             pass
 
         try:
-            endpoint_url = settings.MINIO["ENDPOINT"]
-            access_key = settings.MINIO["ACCESS_KEY"]
-            secret_key = settings.MINIO["SECRET_KEY"]
-            region_name = settings.MINIO["REGION"]
 
             config = Config(
                 s3={
@@ -35,15 +34,15 @@ def __open__(self):
 
             self.conn = boto3.client(
                 's3',
-                endpoint_url=endpoint_url,
-                region_name=region_name,  # 可选，指定区域
-                aws_access_key_id=access_key,
-                aws_secret_access_key=secret_key,
+                endpoint_url=self.endpoint,
+                region_name=self.region,
+                aws_access_key_id=self.access_key,
+                aws_secret_access_key=self.secret_key,
                 config=config
             )
         except Exception as e:
             s3_logger.error(
-                "Fail to connect %s " % settings.MINIO["endpoint_url"] + str(e))
+                "Fail to connect %s " % self.endpoint + str(e))
 
     def __close__(self):
         del self.conn