529 lines
20 KiB
Python
529 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
MongoDB + Redis 数据库缓存管理器
|
||
提供高性能的股票数据缓存和持久化存储
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import pickle
|
||
import hashlib
|
||
from datetime import datetime, timedelta
|
||
from typing import Optional, Dict, Any, List, Union
|
||
import pandas as pd
|
||
|
||
# MongoDB
|
||
try:
|
||
from pymongo import MongoClient
|
||
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
|
||
MONGODB_AVAILABLE = True
|
||
except ImportError:
|
||
MONGODB_AVAILABLE = False
|
||
print("⚠️ pymongo 未安装，MongoDB功能不可用")
|
||
|
||
# Redis
|
||
try:
|
||
import redis
|
||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||
REDIS_AVAILABLE = True
|
||
except ImportError:
|
||
REDIS_AVAILABLE = False
|
||
print("⚠️ redis 未安装，Redis功能不可用")
|
||
|
||
|
||
class DatabaseCacheManager:
|
||
"""MongoDB + Redis 数据库缓存管理器"""
|
||
|
||
def __init__(self,
|
||
mongodb_url: Optional[str] = None,
|
||
redis_url: Optional[str] = None,
|
||
mongodb_db: str = "tradingagents",
|
||
redis_db: int = 0):
|
||
"""
|
||
初始化数据库缓存管理器
|
||
|
||
Args:
|
||
mongodb_url: MongoDB连接URL，默认使用配置文件端口
|
||
redis_url: Redis连接URL，默认使用配置文件端口
|
||
mongodb_db: MongoDB数据库名
|
||
redis_db: Redis数据库编号
|
||
"""
|
||
# 从配置文件获取正确的端口
|
||
mongodb_port = os.getenv("MONGODB_PORT", "27018")
|
||
redis_port = os.getenv("REDIS_PORT", "6380")
|
||
mongodb_password = os.getenv("MONGODB_PASSWORD", "tradingagents123")
|
||
redis_password = os.getenv("REDIS_PASSWORD", "tradingagents123")
|
||
|
||
self.mongodb_url = mongodb_url or os.getenv("MONGODB_URL", f"mongodb://admin:{mongodb_password}@localhost:{mongodb_port}")
|
||
self.redis_url = redis_url or os.getenv("REDIS_URL", f"redis://:{redis_password}@localhost:{redis_port}")
|
||
self.mongodb_db_name = mongodb_db
|
||
self.redis_db = redis_db
|
||
|
||
# 初始化连接
|
||
self.mongodb_client = None
|
||
self.mongodb_db = None
|
||
self.redis_client = None
|
||
|
||
self._init_mongodb()
|
||
self._init_redis()
|
||
|
||
print(f"🗄️ 数据库缓存管理器初始化完成")
|
||
print(f" MongoDB: {'✅ 已连接' if self.mongodb_client else '❌ 未连接'}")
|
||
print(f" Redis: {'✅ 已连接' if self.redis_client else '❌ 未连接'}")
|
||
|
||
def _init_mongodb(self):
|
||
"""初始化MongoDB连接"""
|
||
if not MONGODB_AVAILABLE:
|
||
return
|
||
|
||
try:
|
||
self.mongodb_client = MongoClient(
|
||
self.mongodb_url,
|
||
serverSelectionTimeoutMS=5000, # 5秒超时
|
||
connectTimeoutMS=5000
|
||
)
|
||
# 测试连接
|
||
self.mongodb_client.admin.command('ping')
|
||
self.mongodb_db = self.mongodb_client[self.mongodb_db_name]
|
||
|
||
# 创建索引
|
||
self._create_mongodb_indexes()
|
||
|
||
print(f"✅ MongoDB连接成功: {self.mongodb_url}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ MongoDB连接失败: {e}")
|
||
self.mongodb_client = None
|
||
self.mongodb_db = None
|
||
|
||
def _init_redis(self):
|
||
"""初始化Redis连接"""
|
||
if not REDIS_AVAILABLE:
|
||
return
|
||
|
||
try:
|
||
self.redis_client = redis.from_url(
|
||
self.redis_url,
|
||
db=self.redis_db,
|
||
socket_timeout=5,
|
||
socket_connect_timeout=5,
|
||
decode_responses=True
|
||
)
|
||
# 测试连接
|
||
self.redis_client.ping()
|
||
|
||
print(f"✅ Redis连接成功: {self.redis_url}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ Redis连接失败: {e}")
|
||
self.redis_client = None
|
||
|
||
def _create_mongodb_indexes(self):
|
||
"""创建MongoDB索引"""
|
||
if self.mongodb_db is None:
|
||
return
|
||
|
||
try:
|
||
# 股票数据集合索引
|
||
stock_collection = self.mongodb_db.stock_data
|
||
stock_collection.create_index([
|
||
("symbol", 1),
|
||
("data_source", 1),
|
||
("start_date", 1),
|
||
("end_date", 1)
|
||
])
|
||
stock_collection.create_index([("created_at", 1)])
|
||
|
||
# 新闻数据集合索引
|
||
news_collection = self.mongodb_db.news_data
|
||
news_collection.create_index([
|
||
("symbol", 1),
|
||
("data_source", 1),
|
||
("date_range", 1)
|
||
])
|
||
news_collection.create_index([("created_at", 1)])
|
||
|
||
# 基本面数据集合索引
|
||
fundamentals_collection = self.mongodb_db.fundamentals_data
|
||
fundamentals_collection.create_index([
|
||
("symbol", 1),
|
||
("data_source", 1),
|
||
("analysis_date", 1)
|
||
])
|
||
fundamentals_collection.create_index([("created_at", 1)])
|
||
|
||
print("✅ MongoDB索引创建完成")
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB索引创建失败: {e}")
|
||
|
||
def _generate_cache_key(self, data_type: str, symbol: str, **kwargs) -> str:
|
||
"""生成缓存键"""
|
||
params_str = f"{data_type}_{symbol}"
|
||
for key, value in sorted(kwargs.items()):
|
||
params_str += f"_{key}_{value}"
|
||
|
||
cache_key = hashlib.md5(params_str.encode()).hexdigest()[:16]
|
||
return f"{data_type}:{symbol}:{cache_key}"
|
||
|
||
def save_stock_data(self, symbol: str, data: Union[pd.DataFrame, str],
|
||
start_date: str = None, end_date: str = None,
|
||
data_source: str = "unknown", market_type: str = None) -> str:
|
||
"""
|
||
保存股票数据到MongoDB和Redis
|
||
|
||
Args:
|
||
symbol: 股票代码
|
||
data: 股票数据
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
data_source: 数据源
|
||
market_type: 市场类型 (us/china)
|
||
|
||
Returns:
|
||
cache_key: 缓存键
|
||
"""
|
||
cache_key = self._generate_cache_key("stock", symbol,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
source=data_source)
|
||
|
||
# 自动推断市场类型
|
||
if market_type is None:
|
||
# 根据股票代码格式推断市场类型
|
||
import re
|
||
if re.match(r'^\d{6}$', symbol): # 6位数字为A股
|
||
market_type = "china"
|
||
else: # 其他格式为美股
|
||
market_type = "us"
|
||
|
||
# 准备文档数据
|
||
doc = {
|
||
"_id": cache_key,
|
||
"symbol": symbol,
|
||
"market_type": market_type,
|
||
"data_type": "stock_data",
|
||
"start_date": start_date,
|
||
"end_date": end_date,
|
||
"data_source": data_source,
|
||
"created_at": datetime.utcnow(),
|
||
"updated_at": datetime.utcnow()
|
||
}
|
||
|
||
# 处理数据格式
|
||
if isinstance(data, pd.DataFrame):
|
||
doc["data"] = data.to_json(orient='records', date_format='iso')
|
||
doc["data_format"] = "dataframe_json"
|
||
else:
|
||
doc["data"] = str(data)
|
||
doc["data_format"] = "text"
|
||
|
||
# 保存到MongoDB（持久化）
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
collection = self.mongodb_db.stock_data
|
||
collection.replace_one({"_id": cache_key}, doc, upsert=True)
|
||
print(f"💾 股票数据已保存到MongoDB: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB保存失败: {e}")
|
||
|
||
# 保存到Redis（快速缓存，6小时过期）
|
||
if self.redis_client:
|
||
try:
|
||
redis_data = {
|
||
"data": doc["data"],
|
||
"data_format": doc["data_format"],
|
||
"symbol": symbol,
|
||
"data_source": data_source,
|
||
"created_at": doc["created_at"].isoformat()
|
||
}
|
||
self.redis_client.setex(
|
||
cache_key,
|
||
6 * 3600, # 6小时过期
|
||
json.dumps(redis_data, ensure_ascii=False)
|
||
)
|
||
print(f"⚡ 股票数据已缓存到Redis: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ Redis缓存失败: {e}")
|
||
|
||
return cache_key
|
||
|
||
def load_stock_data(self, cache_key: str) -> Optional[Union[pd.DataFrame, str]]:
|
||
"""从Redis或MongoDB加载股票数据"""
|
||
|
||
# 首先尝试从Redis加载（更快）
|
||
if self.redis_client:
|
||
try:
|
||
redis_data = self.redis_client.get(cache_key)
|
||
if redis_data:
|
||
data_dict = json.loads(redis_data)
|
||
print(f"⚡ 从Redis加载数据: {cache_key}")
|
||
|
||
if data_dict["data_format"] == "dataframe_json":
|
||
return pd.read_json(data_dict["data"], orient='records')
|
||
else:
|
||
return data_dict["data"]
|
||
except Exception as e:
|
||
print(f"⚠️ Redis加载失败: {e}")
|
||
|
||
# 如果Redis没有，从MongoDB加载
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
collection = self.mongodb_db.stock_data
|
||
doc = collection.find_one({"_id": cache_key})
|
||
|
||
if doc:
|
||
print(f"💾 从MongoDB加载数据: {cache_key}")
|
||
|
||
# 同时更新到Redis缓存
|
||
if self.redis_client:
|
||
try:
|
||
redis_data = {
|
||
"data": doc["data"],
|
||
"data_format": doc["data_format"],
|
||
"symbol": doc["symbol"],
|
||
"data_source": doc["data_source"],
|
||
"created_at": doc["created_at"].isoformat()
|
||
}
|
||
self.redis_client.setex(
|
||
cache_key,
|
||
6 * 3600,
|
||
json.dumps(redis_data, ensure_ascii=False)
|
||
)
|
||
print(f"⚡ 数据已同步到Redis缓存")
|
||
except Exception as e:
|
||
print(f"⚠️ Redis同步失败: {e}")
|
||
|
||
if doc["data_format"] == "dataframe_json":
|
||
return pd.read_json(doc["data"], orient='records')
|
||
else:
|
||
return doc["data"]
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB加载失败: {e}")
|
||
|
||
return None
|
||
|
||
def find_cached_stock_data(self, symbol: str, start_date: str = None,
|
||
end_date: str = None, data_source: str = None,
|
||
max_age_hours: int = 6) -> Optional[str]:
|
||
"""查找匹配的缓存数据"""
|
||
|
||
# 生成精确匹配的缓存键
|
||
exact_key = self._generate_cache_key("stock", symbol,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
source=data_source)
|
||
|
||
# 检查Redis中是否有精确匹配
|
||
if self.redis_client and self.redis_client.exists(exact_key):
|
||
print(f"⚡ Redis中找到精确匹配: {symbol} -> {exact_key}")
|
||
return exact_key
|
||
|
||
# 检查MongoDB中的匹配项
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
collection = self.mongodb_db.stock_data
|
||
cutoff_time = datetime.utcnow() - timedelta(hours=max_age_hours)
|
||
|
||
query = {
|
||
"symbol": symbol,
|
||
"created_at": {"$gte": cutoff_time}
|
||
}
|
||
|
||
if data_source:
|
||
query["data_source"] = data_source
|
||
if start_date:
|
||
query["start_date"] = start_date
|
||
if end_date:
|
||
query["end_date"] = end_date
|
||
|
||
doc = collection.find_one(query, sort=[("created_at", -1)])
|
||
|
||
if doc:
|
||
cache_key = doc["_id"]
|
||
print(f"💾 MongoDB中找到匹配: {symbol} -> {cache_key}")
|
||
return cache_key
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB查询失败: {e}")
|
||
|
||
print(f"❌ 未找到有效缓存: {symbol}")
|
||
return None
|
||
|
||
def save_news_data(self, symbol: str, news_data: str,
|
||
start_date: str = None, end_date: str = None,
|
||
data_source: str = "unknown") -> str:
|
||
"""保存新闻数据到MongoDB和Redis"""
|
||
cache_key = self._generate_cache_key("news", symbol,
|
||
start_date=start_date,
|
||
end_date=end_date,
|
||
source=data_source)
|
||
|
||
doc = {
|
||
"_id": cache_key,
|
||
"symbol": symbol,
|
||
"data_type": "news_data",
|
||
"date_range": f"{start_date}_{end_date}",
|
||
"start_date": start_date,
|
||
"end_date": end_date,
|
||
"data_source": data_source,
|
||
"data": news_data,
|
||
"created_at": datetime.utcnow(),
|
||
"updated_at": datetime.utcnow()
|
||
}
|
||
|
||
# 保存到MongoDB
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
collection = self.mongodb_db.news_data
|
||
collection.replace_one({"_id": cache_key}, doc, upsert=True)
|
||
print(f"📰 新闻数据已保存到MongoDB: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB保存失败: {e}")
|
||
|
||
# 保存到Redis（24小时过期）
|
||
if self.redis_client:
|
||
try:
|
||
redis_data = {
|
||
"data": news_data,
|
||
"symbol": symbol,
|
||
"data_source": data_source,
|
||
"created_at": doc["created_at"].isoformat()
|
||
}
|
||
self.redis_client.setex(
|
||
cache_key,
|
||
24 * 3600, # 24小时过期
|
||
json.dumps(redis_data, ensure_ascii=False)
|
||
)
|
||
print(f"⚡ 新闻数据已缓存到Redis: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ Redis缓存失败: {e}")
|
||
|
||
return cache_key
|
||
|
||
def save_fundamentals_data(self, symbol: str, fundamentals_data: str,
|
||
analysis_date: str = None,
|
||
data_source: str = "unknown") -> str:
|
||
"""保存基本面数据到MongoDB和Redis"""
|
||
if not analysis_date:
|
||
analysis_date = datetime.now().strftime("%Y-%m-%d")
|
||
|
||
cache_key = self._generate_cache_key("fundamentals", symbol,
|
||
date=analysis_date,
|
||
source=data_source)
|
||
|
||
doc = {
|
||
"_id": cache_key,
|
||
"symbol": symbol,
|
||
"data_type": "fundamentals_data",
|
||
"analysis_date": analysis_date,
|
||
"data_source": data_source,
|
||
"data": fundamentals_data,
|
||
"created_at": datetime.utcnow(),
|
||
"updated_at": datetime.utcnow()
|
||
}
|
||
|
||
# 保存到MongoDB
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
collection = self.mongodb_db.fundamentals_data
|
||
collection.replace_one({"_id": cache_key}, doc, upsert=True)
|
||
print(f"💼 基本面数据已保存到MongoDB: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB保存失败: {e}")
|
||
|
||
# 保存到Redis（24小时过期）
|
||
if self.redis_client:
|
||
try:
|
||
redis_data = {
|
||
"data": fundamentals_data,
|
||
"symbol": symbol,
|
||
"data_source": data_source,
|
||
"analysis_date": analysis_date,
|
||
"created_at": doc["created_at"].isoformat()
|
||
}
|
||
self.redis_client.setex(
|
||
cache_key,
|
||
24 * 3600, # 24小时过期
|
||
json.dumps(redis_data, ensure_ascii=False)
|
||
)
|
||
print(f"⚡ 基本面数据已缓存到Redis: {symbol} -> {cache_key}")
|
||
except Exception as e:
|
||
print(f"⚠️ Redis缓存失败: {e}")
|
||
|
||
return cache_key
|
||
|
||
def get_cache_stats(self) -> Dict[str, Any]:
|
||
"""获取缓存统计信息"""
|
||
stats = {
|
||
"mongodb": {"available": self.mongodb_db is not None, "collections": {}},
|
||
"redis": {"available": self.redis_client is not None, "keys": 0, "memory_usage": "N/A"}
|
||
}
|
||
|
||
# MongoDB统计
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
for collection_name in ["stock_data", "news_data", "fundamentals_data"]:
|
||
collection = self.mongodb_db[collection_name]
|
||
count = collection.count_documents({})
|
||
size = self.mongodb_db.command("collStats", collection_name).get("size", 0)
|
||
stats["mongodb"]["collections"][collection_name] = {
|
||
"count": count,
|
||
"size_mb": round(size / (1024 * 1024), 2)
|
||
}
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB统计获取失败: {e}")
|
||
|
||
# Redis统计
|
||
if self.redis_client:
|
||
try:
|
||
info = self.redis_client.info()
|
||
stats["redis"]["keys"] = info.get("db0", {}).get("keys", 0)
|
||
stats["redis"]["memory_usage"] = f"{info.get('used_memory_human', 'N/A')}"
|
||
except Exception as e:
|
||
print(f"⚠️ Redis统计获取失败: {e}")
|
||
|
||
return stats
|
||
|
||
def clear_old_cache(self, max_age_days: int = 7):
|
||
"""清理过期缓存"""
|
||
cutoff_time = datetime.utcnow() - timedelta(days=max_age_days)
|
||
cleared_count = 0
|
||
|
||
# 清理MongoDB
|
||
if self.mongodb_db is not None:
|
||
try:
|
||
for collection_name in ["stock_data", "news_data", "fundamentals_data"]:
|
||
collection = self.mongodb_db[collection_name]
|
||
result = collection.delete_many({"created_at": {"$lt": cutoff_time}})
|
||
cleared_count += result.deleted_count
|
||
print(f"🧹 MongoDB {collection_name} 清理了 {result.deleted_count} 条记录")
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB清理失败: {e}")
|
||
|
||
# Redis会自动过期，不需要手动清理
|
||
print(f"🧹 总共清理了 {cleared_count} 条过期记录")
|
||
return cleared_count
|
||
|
||
def close(self):
|
||
"""关闭数据库连接"""
|
||
if self.mongodb_client:
|
||
self.mongodb_client.close()
|
||
print("🔒 MongoDB连接已关闭")
|
||
|
||
if self.redis_client:
|
||
self.redis_client.close()
|
||
print("🔒 Redis连接已关闭")
|
||
|
||
|
||
# 全局数据库缓存实例
|
||
_db_cache_instance = None
|
||
|
||
def get_db_cache() -> DatabaseCacheManager:
|
||
"""获取全局数据库缓存实例"""
|
||
global _db_cache_instance
|
||
if _db_cache_instance is None:
|
||
_db_cache_instance = DatabaseCacheManager()
|
||
return _db_cache_instance
|