refactor(database): 重新调整docker容器的连接配置

management/server/database.py

@@ -1,7 +1,5 @@
 import mysql.connector
 import os
-from utils import generate_uuid, encrypt_password
-from datetime import datetime
 from minio import Minio
 from dotenv import load_dotenv
 from elasticsearch import Elasticsearch
@@ -9,26 +7,39 @@ from elasticsearch import Elasticsearch
 # 加载环境变量
 load_dotenv("../../docker/.env")
 
+
 # 检测是否在Docker容器中运行
 def is_running_in_docker():
     # 检查是否存在/.dockerenv文件
-    docker_env = os.path.exists('/.dockerenv')
+    docker_env = os.path.exists("/.dockerenv")
     # 或者检查cgroup中是否包含docker字符串
     try:
-        with open('/proc/self/cgroup', 'r') as f:
-            return docker_env or 'docker' in f.read()
-    except:
+        with open("/proc/self/cgroup", "r") as f:
+            return docker_env or "docker" in f.read()
+    except:  # noqa: E722
         return docker_env
 
-# 根据运行环境选择合适的主机地址
-DB_HOST = 'host.docker.internal' if is_running_in_docker() else 'localhost'
-MINIO_HOST = 'host.docker.internal' if is_running_in_docker() else 'localhost'
-ES_HOST = 'es01' if is_running_in_docker() else 'localhost'
+
+# 根据运行环境选择合适的主机地址和端口
+if is_running_in_docker():
+    MYSQL_HOST = "mysql"
+    MYSQL_PORT = 3306
+    MINIO_HOST = "minio"
+    MINIO_PORT = 9000
+    ES_HOST = "es01"
+    ES_PORT = 9200
+else:
+    MYSQL_HOST = "host.docker.internal"
+    MYSQL_PORT = int(os.getenv("MYSQL_PORT", "5455"))
+    MINIO_HOST = "host.docker.internal"
+    MINIO_PORT = int(os.getenv("MINIO_PORT", "9000"))
+    ES_HOST = "host.docker.internal"
+    ES_PORT = int(os.getenv("ES_PORT", "9200"))
 
 # 数据库连接配置
 DB_CONFIG = {
-    "host": DB_HOST,
-    "port": int(os.getenv("MYSQL_PORT", "5455")),
+    "host": MYSQL_HOST,
+    "port": MYSQL_PORT,
     "user": "root",
     "password": os.getenv("MYSQL_PASSWORD", "infini_rag_flow"),
     "database": "rag_flow",
@@ -36,20 +47,21 @@ DB_CONFIG = {
 
 # MinIO连接配置
 MINIO_CONFIG = {
-    "endpoint": f"{MINIO_HOST}:{os.getenv('MINIO_PORT', '9000')}",
+    "endpoint": f"{MINIO_HOST}:{MINIO_PORT}",
     "access_key": os.getenv("MINIO_USER", "rag_flow"),
     "secret_key": os.getenv("MINIO_PASSWORD", "infini_rag_flow"),
-    "secure": False
+    "secure": False,
 }
 
 # Elasticsearch连接配置
 ES_CONFIG = {
-    "host": f"http://{ES_HOST}:{os.getenv('ES_PORT', '9200')}", 
+    "host": f"http://{ES_HOST}:{ES_PORT}",
     "user": os.getenv("ELASTIC_USER", "elastic"),
     "password": os.getenv("ELASTIC_PASSWORD", "infini_rag_flow"),
-    "use_ssl": os.getenv("ES_USE_SSL", "false").lower() == "true"
+    "use_ssl": os.getenv("ES_USE_SSL", "false").lower() == "true",
 }
 
+
 def get_db_connection():
     """创建MySQL数据库连接"""
     try:
@@ -59,27 +71,22 @@ def get_db_connection():
         print(f"MySQL连接失败: {str(e)}")
         raise e
 
+
 def get_minio_client():
     """创建MinIO客户端连接"""
     try:
-        minio_client = Minio(
-            endpoint=MINIO_CONFIG["endpoint"],
-            access_key=MINIO_CONFIG["access_key"],
-            secret_key=MINIO_CONFIG["secret_key"],
-            secure=MINIO_CONFIG["secure"]
-        )
+        minio_client = Minio(endpoint=MINIO_CONFIG["endpoint"], access_key=MINIO_CONFIG["access_key"], secret_key=MINIO_CONFIG["secret_key"], secure=MINIO_CONFIG["secure"])
         return minio_client
     except Exception as e:
         print(f"MinIO连接失败: {str(e)}")
         raise e
 
+
 def get_es_client():
     """创建Elasticsearch客户端连接"""
     try:
         # 构建连接参数
-        es_params = {
-            "hosts": [ES_CONFIG["host"]]
-        }
+        es_params = {"hosts": [ES_CONFIG["host"]]}
 
         # 如果提供了用户名和密码，添加认证信息
         if ES_CONFIG["user"] and ES_CONFIG["password"]:
@@ -96,6 +103,7 @@ def get_es_client():
         print(f"Elasticsearch连接失败: {str(e)}")
         raise e
 
+
 def test_connections():
     """测试数据库和MinIO连接"""
     try:
@@ -126,6 +134,7 @@ def test_connections():
         print(f"连接测试失败: {str(e)}")
         return False
 
+
 if __name__ == "__main__":
     # 如果直接运行此文件，则测试连接
     test_connections()