浏览代码

refactor(数据库配置): 重构数据库连接配置以支持环境变量和本地开发

从代码中移除硬编码的数据库配置,改为通过环境变量加载
添加对本地开发环境的支持,使用不同的测试数据库
统一两个文件中的数据库配置逻辑
max_liu 2 天之前
父节点
当前提交
efae2334d9
共有 2 个文件被更改,包括 20 次插入10 次删除
  1. 5 3
      src/models/database.py
  2. 15 7
      utils/sync_mysql_help.py

+ 5 - 3
src/models/database.py

@@ -5,6 +5,7 @@
 """
 
 import os
+from dotenv import load_dotenv
 from typing import Generator
 from sqlalchemy import create_engine
 from sqlalchemy.ext.declarative import declarative_base
@@ -36,11 +37,13 @@ def get_database_url() -> str:
         DB_PASSWORD: 数据库密码 (必需)
         DB_NAME: 数据库名称 (默认: content-deconstruction)
     """
+    load_dotenv()
+    env = (os.getenv("APP_ENV") or os.getenv("ENV") or "local").lower()
     host = os.getenv("DB_HOST", "rm-t4nh1xx6o2a6vj8qu3o.mysql.singapore.rds.aliyuncs.com")
     port = os.getenv("DB_PORT", "3306")
     user = os.getenv("DB_USER", "content_rw")
     password = os.getenv("DB_PASSWORD", "bC1aH4bA1lB0")
-    database = os.getenv("DB_NAME", "content-deconstruction")
+    database = os.getenv("DB_NAME", "content-deconstruction-test" if env in ("local","dev","development") else "content-deconstruction")
     
     if not password:
         raise ValueError("DB_PASSWORD environment variable is required")
@@ -65,7 +68,7 @@ def get_engine():
             pool_pre_ping=True,  # 连接前检查连接是否有效
             echo=False,  # 设置为 True 可以打印 SQL 语句,用于调试
         )
-        logger.info(f"Database engine created for database: {os.getenv('DB_NAME', 'content-deconstruction')}")
+        logger.info(f"Database engine created for database: {os.getenv('DB_NAME', database)}")
     return _engine
 
 
@@ -123,4 +126,3 @@ def drop_db():
     engine = get_engine()
     Base.metadata.drop_all(bind=engine)
     logger.warning("All database tables dropped")
-

+ 15 - 7
utils/sync_mysql_help.py

@@ -1,4 +1,6 @@
+import os
 import pymysql
+from dotenv import load_dotenv
 
 
 from typing import Tuple, Any, Dict, Literal, Optional
@@ -19,16 +21,24 @@ class SyncMySQLHelper(object):
 
     def get_pool(self):
         if self._pool is None:
+            # 加载环境变量,允许通过 .env 配置本机调试数据库
+            load_dotenv()
+            env = (os.getenv('APP_ENV') or os.getenv('ENV') or 'local').lower()
+            host = os.getenv('DB_HOST', 'rm-t4nh1xx6o2a6vj8qu3o.mysql.singapore.rds.aliyuncs.com')
+            port = int(os.getenv('DB_PORT', '3306'))
+            user = os.getenv('DB_USER', 'content_rw')
+            password = os.getenv('DB_PASSWORD', 'bC1aH4bA1lB0')
+            database = os.getenv('DB_NAME', 'content-deconstruction-test' if env in ('local','dev','development') else 'content-deconstruction')
             self._pool = PooledDB(
                 creator=pymysql,
                 mincached=10,
                 maxconnections=20,
                 blocking=True,
-                host='rm-t4nh1xx6o2a6vj8qu3o.mysql.singapore.rds.aliyuncs.com',
-                port=3306,
-                user='content_rw',
-                password='bC1aH4bA1lB0',
-                database='content-deconstruction')
+                host=host,
+                port=port,
+                user=user,
+                password=password,
+                database=database)
 
         return self._pool
 
@@ -82,5 +92,3 @@ class SyncMySQLHelper(object):
 
 mysql = SyncMySQLHelper()
 
-
-