361 lines
12 KiB
Python
361 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
智能数据库管理器
|
||
自动检测MongoDB和Redis可用性，提供降级方案
|
||
使用项目现有的.env配置
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Dict, Any, Optional, Tuple
|
||
|
||
class DatabaseManager:
|
||
"""智能数据库管理器"""
|
||
|
||
def __init__(self):
|
||
self.logger = logging.getLogger(__name__)
|
||
|
||
# 加载.env配置
|
||
self._load_env_config()
|
||
|
||
# 数据库连接状态
|
||
self.mongodb_available = False
|
||
self.redis_available = False
|
||
self.mongodb_client = None
|
||
self.redis_client = None
|
||
|
||
# 检测数据库可用性
|
||
self._detect_databases()
|
||
|
||
# 初始化连接
|
||
self._initialize_connections()
|
||
|
||
self.logger.info(f"数据库管理器初始化完成 - MongoDB: {self.mongodb_available}, Redis: {self.redis_available}")
|
||
|
||
def _load_env_config(self):
|
||
"""从.env文件加载配置"""
|
||
# 尝试加载python-dotenv
|
||
try:
|
||
from dotenv import load_dotenv
|
||
load_dotenv()
|
||
except ImportError:
|
||
self.logger.info("python-dotenv未安装，直接读取环境变量")
|
||
|
||
# 读取启用开关
|
||
self.mongodb_enabled = os.getenv("MONGODB_ENABLED", "false").lower() == "true"
|
||
self.redis_enabled = os.getenv("REDIS_ENABLED", "false").lower() == "true"
|
||
|
||
# 从环境变量读取MongoDB配置
|
||
self.mongodb_config = {
|
||
"enabled": self.mongodb_enabled,
|
||
"host": os.getenv("MONGODB_HOST", "localhost"),
|
||
"port": int(os.getenv("MONGODB_PORT", "27017")),
|
||
"username": os.getenv("MONGODB_USERNAME"),
|
||
"password": os.getenv("MONGODB_PASSWORD"),
|
||
"database": os.getenv("MONGODB_DATABASE", "tradingagents"),
|
||
"auth_source": os.getenv("MONGODB_AUTH_SOURCE", "admin"),
|
||
"timeout": 2000
|
||
}
|
||
|
||
# 从环境变量读取Redis配置
|
||
self.redis_config = {
|
||
"enabled": self.redis_enabled,
|
||
"host": os.getenv("REDIS_HOST", "localhost"),
|
||
"port": int(os.getenv("REDIS_PORT", "6379")),
|
||
"password": os.getenv("REDIS_PASSWORD"),
|
||
"db": int(os.getenv("REDIS_DB", "0")),
|
||
"timeout": 2
|
||
}
|
||
|
||
self.logger.info(f"MongoDB启用: {self.mongodb_enabled}")
|
||
self.logger.info(f"Redis启用: {self.redis_enabled}")
|
||
if self.mongodb_enabled:
|
||
self.logger.info(f"MongoDB配置: {self.mongodb_config['host']}:{self.mongodb_config['port']}")
|
||
if self.redis_enabled:
|
||
self.logger.info(f"Redis配置: {self.redis_config['host']}:{self.redis_config['port']}")
|
||
|
||
|
||
|
||
def _detect_mongodb(self) -> Tuple[bool, str]:
|
||
"""检测MongoDB是否可用"""
|
||
# 首先检查是否启用
|
||
if not self.mongodb_enabled:
|
||
return False, "MongoDB未启用 (MONGODB_ENABLED=false)"
|
||
|
||
try:
|
||
import pymongo
|
||
from pymongo import MongoClient
|
||
|
||
# 构建连接参数
|
||
connect_kwargs = {
|
||
"host": self.mongodb_config["host"],
|
||
"port": self.mongodb_config["port"],
|
||
"serverSelectionTimeoutMS": self.mongodb_config["timeout"],
|
||
"connectTimeoutMS": self.mongodb_config["timeout"]
|
||
}
|
||
|
||
# 如果有用户名和密码，添加认证
|
||
if self.mongodb_config["username"] and self.mongodb_config["password"]:
|
||
connect_kwargs.update({
|
||
"username": self.mongodb_config["username"],
|
||
"password": self.mongodb_config["password"],
|
||
"authSource": self.mongodb_config["auth_source"]
|
||
})
|
||
|
||
client = MongoClient(**connect_kwargs)
|
||
|
||
# 测试连接
|
||
client.server_info()
|
||
client.close()
|
||
|
||
return True, "MongoDB连接成功"
|
||
|
||
except ImportError:
|
||
return False, "pymongo未安装"
|
||
except Exception as e:
|
||
return False, f"MongoDB连接失败: {str(e)}"
|
||
|
||
def _detect_redis(self) -> Tuple[bool, str]:
|
||
"""检测Redis是否可用"""
|
||
# 首先检查是否启用
|
||
if not self.redis_enabled:
|
||
return False, "Redis未启用 (REDIS_ENABLED=false)"
|
||
|
||
try:
|
||
import redis
|
||
|
||
# 构建连接参数
|
||
connect_kwargs = {
|
||
"host": self.redis_config["host"],
|
||
"port": self.redis_config["port"],
|
||
"db": self.redis_config["db"],
|
||
"socket_timeout": self.redis_config["timeout"],
|
||
"socket_connect_timeout": self.redis_config["timeout"]
|
||
}
|
||
|
||
# 如果有密码，添加密码
|
||
if self.redis_config["password"]:
|
||
connect_kwargs["password"] = self.redis_config["password"]
|
||
|
||
client = redis.Redis(**connect_kwargs)
|
||
|
||
# 测试连接
|
||
client.ping()
|
||
|
||
return True, "Redis连接成功"
|
||
|
||
except ImportError:
|
||
return False, "redis未安装"
|
||
except Exception as e:
|
||
return False, f"Redis连接失败: {str(e)}"
|
||
|
||
def _detect_databases(self):
|
||
"""检测所有数据库"""
|
||
self.logger.info("开始检测数据库可用性...")
|
||
|
||
# 检测MongoDB
|
||
mongodb_available, mongodb_msg = self._detect_mongodb()
|
||
self.mongodb_available = mongodb_available
|
||
|
||
if mongodb_available:
|
||
self.logger.info(f"✅ MongoDB: {mongodb_msg}")
|
||
else:
|
||
self.logger.info(f"❌ MongoDB: {mongodb_msg}")
|
||
|
||
# 检测Redis
|
||
redis_available, redis_msg = self._detect_redis()
|
||
self.redis_available = redis_available
|
||
|
||
if redis_available:
|
||
self.logger.info(f"✅ Redis: {redis_msg}")
|
||
else:
|
||
self.logger.info(f"❌ Redis: {redis_msg}")
|
||
|
||
# 更新配置
|
||
self._update_config_based_on_detection()
|
||
|
||
def _update_config_based_on_detection(self):
|
||
"""根据检测结果更新配置"""
|
||
# 确定缓存后端
|
||
if self.redis_available:
|
||
self.primary_backend = "redis"
|
||
elif self.mongodb_available:
|
||
self.primary_backend = "mongodb"
|
||
else:
|
||
self.primary_backend = "file"
|
||
|
||
self.logger.info(f"主要缓存后端: {self.primary_backend}")
|
||
|
||
def _initialize_connections(self):
|
||
"""初始化数据库连接"""
|
||
# 初始化MongoDB连接
|
||
if self.mongodb_available:
|
||
try:
|
||
import pymongo
|
||
|
||
# 构建连接参数
|
||
connect_kwargs = {
|
||
"host": self.mongodb_config["host"],
|
||
"port": self.mongodb_config["port"],
|
||
"serverSelectionTimeoutMS": self.mongodb_config["timeout"]
|
||
}
|
||
|
||
# 如果有用户名和密码，添加认证
|
||
if self.mongodb_config["username"] and self.mongodb_config["password"]:
|
||
connect_kwargs.update({
|
||
"username": self.mongodb_config["username"],
|
||
"password": self.mongodb_config["password"],
|
||
"authSource": self.mongodb_config["auth_source"]
|
||
})
|
||
|
||
self.mongodb_client = pymongo.MongoClient(**connect_kwargs)
|
||
self.logger.info("MongoDB客户端初始化成功")
|
||
except Exception as e:
|
||
self.logger.error(f"MongoDB客户端初始化失败: {e}")
|
||
self.mongodb_available = False
|
||
|
||
# 初始化Redis连接
|
||
if self.redis_available:
|
||
try:
|
||
import redis
|
||
|
||
# 构建连接参数
|
||
connect_kwargs = {
|
||
"host": self.redis_config["host"],
|
||
"port": self.redis_config["port"],
|
||
"db": self.redis_config["db"],
|
||
"socket_timeout": self.redis_config["timeout"]
|
||
}
|
||
|
||
# 如果有密码，添加密码
|
||
if self.redis_config["password"]:
|
||
connect_kwargs["password"] = self.redis_config["password"]
|
||
|
||
self.redis_client = redis.Redis(**connect_kwargs)
|
||
self.logger.info("Redis客户端初始化成功")
|
||
except Exception as e:
|
||
self.logger.error(f"Redis客户端初始化失败: {e}")
|
||
self.redis_available = False
|
||
|
||
def get_mongodb_client(self):
|
||
"""获取MongoDB客户端"""
|
||
if self.mongodb_available and self.mongodb_client:
|
||
return self.mongodb_client
|
||
return None
|
||
|
||
def get_redis_client(self):
|
||
"""获取Redis客户端"""
|
||
if self.redis_available and self.redis_client:
|
||
return self.redis_client
|
||
return None
|
||
|
||
def is_mongodb_available(self) -> bool:
|
||
"""检查MongoDB是否可用"""
|
||
return self.mongodb_available
|
||
|
||
def is_redis_available(self) -> bool:
|
||
"""检查Redis是否可用"""
|
||
return self.redis_available
|
||
|
||
def is_database_available(self) -> bool:
|
||
"""检查是否有任何数据库可用"""
|
||
return self.mongodb_available or self.redis_available
|
||
|
||
def get_cache_backend(self) -> str:
|
||
"""获取当前缓存后端"""
|
||
return self.primary_backend
|
||
|
||
def get_config(self) -> Dict[str, Any]:
|
||
"""获取配置信息"""
|
||
return {
|
||
"mongodb": self.mongodb_config,
|
||
"redis": self.redis_config,
|
||
"primary_backend": self.primary_backend,
|
||
"mongodb_available": self.mongodb_available,
|
||
"redis_available": self.redis_available
|
||
}
|
||
|
||
def get_status_report(self) -> Dict[str, Any]:
|
||
"""获取状态报告"""
|
||
return {
|
||
"database_available": self.is_database_available(),
|
||
"mongodb": {
|
||
"available": self.mongodb_available,
|
||
"host": self.mongodb_config["host"],
|
||
"port": self.mongodb_config["port"]
|
||
},
|
||
"redis": {
|
||
"available": self.redis_available,
|
||
"host": self.redis_config["host"],
|
||
"port": self.redis_config["port"]
|
||
},
|
||
"cache_backend": self.get_cache_backend(),
|
||
"fallback_enabled": True # 总是启用降级
|
||
}
|
||
|
||
def get_cache_stats(self) -> Dict[str, Any]:
|
||
"""获取缓存统计信息"""
|
||
stats = {
|
||
"mongodb_available": self.mongodb_available,
|
||
"redis_available": self.redis_available,
|
||
"redis_keys": 0,
|
||
"redis_memory": "N/A"
|
||
}
|
||
|
||
# Redis统计
|
||
if self.redis_available and self.redis_client:
|
||
try:
|
||
info = self.redis_client.info()
|
||
stats["redis_keys"] = self.redis_client.dbsize()
|
||
stats["redis_memory"] = info.get("used_memory_human", "N/A")
|
||
except Exception as e:
|
||
self.logger.error(f"获取Redis统计失败: {e}")
|
||
|
||
return stats
|
||
|
||
def cache_clear_pattern(self, pattern: str) -> int:
|
||
"""清理匹配模式的缓存"""
|
||
cleared_count = 0
|
||
|
||
if self.redis_available and self.redis_client:
|
||
try:
|
||
keys = self.redis_client.keys(pattern)
|
||
if keys:
|
||
cleared_count += self.redis_client.delete(*keys)
|
||
except Exception as e:
|
||
self.logger.error(f"Redis缓存清理失败: {e}")
|
||
|
||
return cleared_count
|
||
|
||
|
||
# 全局数据库管理器实例
|
||
_database_manager = None
|
||
|
||
def get_database_manager() -> DatabaseManager:
|
||
"""获取全局数据库管理器实例"""
|
||
global _database_manager
|
||
if _database_manager is None:
|
||
_database_manager = DatabaseManager()
|
||
return _database_manager
|
||
|
||
def is_mongodb_available() -> bool:
|
||
"""检查MongoDB是否可用"""
|
||
return get_database_manager().is_mongodb_available()
|
||
|
||
def is_redis_available() -> bool:
|
||
"""检查Redis是否可用"""
|
||
return get_database_manager().is_redis_available()
|
||
|
||
def get_cache_backend() -> str:
|
||
"""获取当前缓存后端"""
|
||
return get_database_manager().get_cache_backend()
|
||
|
||
def get_mongodb_client():
|
||
"""获取MongoDB客户端"""
|
||
return get_database_manager().get_mongodb_client()
|
||
|
||
def get_redis_client():
|
||
"""获取Redis客户端"""
|
||
return get_database_manager().get_redis_client()
|