651 lines
22 KiB
Python
651 lines
22 KiB
Python
"""
|
|
PostgreSQL存储模块
|
|
实现公告数据的PostgreSQL存储和管理
|
|
"""
|
|
|
|
from typing import List, Dict, Any, Optional
|
|
from datetime import datetime, timedelta
|
|
import threading
|
|
from dataclasses import asdict
|
|
|
|
try:
|
|
from ..core.models import Announcement, CrawlResult
|
|
from ..core.database import get_database_manager, init_database
|
|
from ..core.logger import get_logger
|
|
from ..core.reliability import retry_on_exception, RetryConfig, safe_execute
|
|
except ImportError:
|
|
from core.models import Announcement, CrawlResult
|
|
from core.database import get_database_manager, init_database
|
|
from core.logger import get_logger
|
|
from core.reliability import retry_on_exception, RetryConfig, safe_execute
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class PostgreSQLStorage:
|
|
"""PostgreSQL存储管理器"""
|
|
|
|
def __init__(self):
|
|
self.db_manager = get_database_manager()
|
|
self._lock = threading.Lock()
|
|
|
|
def save_announcements(self, announcements: List[Announcement]) -> int:
|
|
"""
|
|
保存公告列表到数据库（经过筛选的公告）
|
|
|
|
Args:
|
|
announcements: 公告列表
|
|
|
|
Returns:
|
|
int: 成功保存的数量
|
|
"""
|
|
if not announcements:
|
|
return 0
|
|
|
|
logger.info(f"开始保存 {len(announcements)} 条筛选后公告到数据库")
|
|
|
|
try:
|
|
# 批量保存
|
|
saved_count = self.db_manager.save_announcements_batch(announcements)
|
|
|
|
if saved_count > 0:
|
|
logger.info(f"成功保存 {saved_count} 条筛选后公告到数据库")
|
|
|
|
# 标记新公告
|
|
self._mark_new_announcements(announcements[:saved_count])
|
|
|
|
return saved_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"保存筛选后公告到数据库失败: {str(e)}")
|
|
# 尝试逐个保存
|
|
return self._save_announcements_fallback(announcements)
|
|
|
|
def save_auto_announcements(self, announcements: List[Announcement]) -> int:
|
|
"""
|
|
保存定时搜索公告到专用表（关键词匹配专用）
|
|
|
|
Args:
|
|
announcements: 公告列表
|
|
|
|
Returns:
|
|
int: 成功保存的数量
|
|
"""
|
|
if not announcements:
|
|
return 0
|
|
|
|
logger.info(f"开始保存 {len(announcements)} 条定时搜索公告到专用表")
|
|
|
|
try:
|
|
# 批量保存到auto_announcements表
|
|
saved_count = self.db_manager.save_announcements_batch_to_table(announcements, "auto_announcements")
|
|
|
|
if saved_count > 0:
|
|
logger.info(f"成功保存 {saved_count} 条定时搜索公告到专用表")
|
|
|
|
return saved_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"保存定时搜索公告到专用表失败: {str(e)}")
|
|
return 0
|
|
|
|
def save_manual_announcements_by_source(self, announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""
|
|
按来源保存手动搜索公告到专用表，每个来源保留最新的max_per_source条
|
|
|
|
Args:
|
|
announcements: 所有公告列表（未经关键词筛选）
|
|
max_per_source: 每个来源最大保留数量
|
|
|
|
Returns:
|
|
Dict[str, int]: 各来源保存的数量
|
|
"""
|
|
if not announcements:
|
|
return {}
|
|
|
|
logger.info(f"开始按来源保存 {len(announcements)} 条手动搜索公告到专用表，每个来源最多保留 {max_per_source} 条")
|
|
|
|
try:
|
|
# 按来源分组
|
|
source_groups = {}
|
|
for announcement in announcements:
|
|
source_code = announcement.source_code
|
|
if source_code not in source_groups:
|
|
source_groups[source_code] = []
|
|
source_groups[source_code].append(announcement)
|
|
|
|
saved_stats = {}
|
|
|
|
for source_code, source_announcements in source_groups.items():
|
|
# 对每个来源的公告按发布时间排序（最新的在前）
|
|
sorted_announcements = sorted(
|
|
source_announcements,
|
|
key=lambda x: x.publish_date or datetime.min,
|
|
reverse=True
|
|
)
|
|
|
|
# 为没有哈希的公告生成哈希
|
|
for announcement in sorted_announcements:
|
|
if not announcement.content_hash:
|
|
announcement.generate_content_hash()
|
|
|
|
# 批量保存
|
|
to_save = sorted_announcements[:max_per_source]
|
|
saved_count = self.db_manager.save_announcements_batch_to_table(to_save, "manual_announcements")
|
|
saved_stats[source_code] = saved_count
|
|
|
|
# 清理该来源超出限制的旧数据
|
|
if len(sorted_announcements) > max_per_source:
|
|
self._cleanup_old_announcements_by_source_in_table(source_code, max_per_source, "manual_announcements")
|
|
|
|
logger.info(f"来源 {source_code} 保存了 {saved_count} 条手动搜索公告")
|
|
|
|
total_saved = sum(saved_stats.values())
|
|
logger.info(f"按来源保存手动搜索公告完成，总计保存 {total_saved} 条公告")
|
|
|
|
return saved_stats
|
|
|
|
except Exception as e:
|
|
logger.error(f"按来源保存手动搜索公告失败: {str(e)}")
|
|
return {}
|
|
|
|
def save_all_announcements_by_source(self, announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""
|
|
按来源保存所有公告，每个来源保留最新的max_per_source条
|
|
|
|
Args:
|
|
announcements: 所有公告列表（未经关键词筛选）
|
|
max_per_source: 每个来源最大保留数量
|
|
|
|
Returns:
|
|
Dict[str, int]: 各来源保存的数量
|
|
"""
|
|
if not announcements:
|
|
return {}
|
|
|
|
logger.info(f"开始按来源保存 {len(announcements)} 条公告，每个来源最多保留 {max_per_source} 条")
|
|
|
|
try:
|
|
# 按来源分组
|
|
source_groups = {}
|
|
for announcement in announcements:
|
|
source_code = announcement.source_code
|
|
if source_code not in source_groups:
|
|
source_groups[source_code] = []
|
|
source_groups[source_code].append(announcement)
|
|
|
|
saved_stats = {}
|
|
|
|
for source_code, source_announcements in source_groups.items():
|
|
# 对每个来源的公告按发布时间排序（最新的在前）
|
|
sorted_announcements = sorted(
|
|
source_announcements,
|
|
key=lambda x: x.publish_date or x.crawled_at or datetime.min,
|
|
reverse=True
|
|
)
|
|
|
|
# 取最新的max_per_source条
|
|
to_save = sorted_announcements[:max_per_source]
|
|
|
|
# 为这些公告生成哈希
|
|
for announcement in to_save:
|
|
if not announcement.content_hash:
|
|
announcement.generate_content_hash()
|
|
|
|
# 批量保存
|
|
saved_count = self.db_manager.save_announcements_batch(to_save)
|
|
|
|
saved_stats[source_code] = saved_count
|
|
|
|
# 清理该来源超出限制的旧数据
|
|
if len(sorted_announcements) > max_per_source:
|
|
self._cleanup_old_announcements_by_source(source_code, max_per_source)
|
|
|
|
logger.info(f"来源 {source_code} 保存了 {saved_count} 条公告")
|
|
|
|
total_saved = sum(saved_stats.values())
|
|
logger.info(f"按来源保存完成，总计保存 {total_saved} 条公告")
|
|
|
|
return saved_stats
|
|
|
|
except Exception as e:
|
|
logger.error(f"按来源保存公告失败: {str(e)}")
|
|
return {}
|
|
|
|
def _cleanup_old_announcements_by_source(self, source_code: str, keep_count: int):
|
|
"""
|
|
清理指定来源超出限制的旧公告
|
|
|
|
Args:
|
|
source_code: 来源代码
|
|
keep_count: 保留数量
|
|
"""
|
|
try:
|
|
# 使用窗口函数删除超出限制的记录
|
|
sql = """
|
|
DELETE FROM announcements
|
|
WHERE source_code = %s
|
|
AND id IN (
|
|
SELECT id FROM (
|
|
SELECT id,
|
|
ROW_NUMBER() OVER (ORDER BY publish_date DESC, crawled_at DESC) as rn
|
|
FROM announcements
|
|
WHERE source_code = %s
|
|
) ranked
|
|
WHERE rn > %s
|
|
)
|
|
"""
|
|
|
|
with get_db_cursor() as cursor:
|
|
cursor.execute(sql, (source_code, source_code, keep_count))
|
|
deleted_count = cursor.rowcount
|
|
|
|
if deleted_count > 0:
|
|
logger.debug(f"清理来源 {source_code} 的 {deleted_count} 条旧公告")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"清理来源 {source_code} 旧公告失败: {str(e)}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"按来源保存公告失败: {str(e)}")
|
|
return {}
|
|
|
|
def _save_announcements_fallback(self, announcements: List[Announcement]) -> int:
|
|
"""逐个保存公告的降级方案"""
|
|
logger.info("使用降级方案逐个保存公告")
|
|
|
|
saved_count = 0
|
|
for announcement in announcements:
|
|
try:
|
|
if self.db_manager.save_announcement(announcement):
|
|
saved_count += 1
|
|
except Exception as e:
|
|
logger.warning(f"保存公告失败: {announcement.title[:50]}..., 错误: {str(e)}")
|
|
continue
|
|
|
|
logger.info(f"降级保存完成，成功保存 {saved_count} 条公告")
|
|
return saved_count
|
|
|
|
def _mark_new_announcements(self, announcements: List[Announcement]):
|
|
"""标记新公告"""
|
|
# 这里可以添加新公告标记逻辑
|
|
# 由于我们在爬取时已经标记，这里主要是确保数据库中的标记正确
|
|
pass
|
|
|
|
def _mark_new_announcements_in_table(self, announcements: List[Announcement], table_name: str):
|
|
"""在指定表中标记新公告"""
|
|
# 这里可以添加新公告标记逻辑
|
|
pass
|
|
|
|
def _cleanup_old_announcements_by_source_in_table(self, source_code: str, max_per_source: int, table_name: str):
|
|
"""在指定表中清理来源的旧公告"""
|
|
try:
|
|
with self.db_manager.get_db_cursor() as cursor:
|
|
# 获取该来源当前保存的公告数量
|
|
cursor.execute(f"""
|
|
SELECT COUNT(*) FROM {table_name}
|
|
WHERE source_code = %s
|
|
""", (source_code,))
|
|
|
|
current_count = cursor.fetchone()[0]
|
|
|
|
if current_count > max_per_source:
|
|
# 删除超出数量的旧公告
|
|
delete_count = current_count - max_per_source
|
|
cursor.execute(f"""
|
|
DELETE FROM {table_name}
|
|
WHERE id IN (
|
|
SELECT id FROM {table_name}
|
|
WHERE source_code = %s
|
|
ORDER BY publish_date DESC, created_at DESC
|
|
OFFSET %s
|
|
)
|
|
""", (source_code, max_per_source))
|
|
|
|
logger.info(f"清理了 {cursor.rowcount} 条{table_name}表中来源{source_code}的旧公告")
|
|
|
|
except Exception as e:
|
|
logger.error(f"清理{table_name}表中来源{source_code}的旧公告失败: {str(e)}")
|
|
|
|
def save_crawl_results(self, results: List[CrawlResult]) -> int:
|
|
"""
|
|
保存爬取结果
|
|
|
|
Args:
|
|
results: 爬取结果列表
|
|
|
|
Returns:
|
|
int: 成功保存的数量
|
|
"""
|
|
if not results:
|
|
return 0
|
|
|
|
saved_count = 0
|
|
for result in results:
|
|
try:
|
|
if self.db_manager.save_crawl_result(result):
|
|
saved_count += 1
|
|
except Exception as e:
|
|
logger.warning(f"保存爬取结果失败: {result.source.name}, 错误: {str(e)}")
|
|
continue
|
|
|
|
logger.info(f"保存爬取结果完成: {saved_count}/{len(results)}")
|
|
return saved_count
|
|
|
|
def get_recent_announcements(self, hours: int = 24,
|
|
limit: int = 100) -> List[Announcement]:
|
|
"""
|
|
获取最近的公告
|
|
|
|
Args:
|
|
hours: 最近小时数
|
|
limit: 限制数量
|
|
|
|
Returns:
|
|
List[Announcement]: 公告列表
|
|
"""
|
|
try:
|
|
return self.db_manager.get_recent_announcements(hours)
|
|
except Exception as e:
|
|
logger.error(f"获取最近公告失败: {str(e)}")
|
|
return []
|
|
|
|
def cleanup_expired_data(self, days: Optional[int] = None) -> int:
|
|
"""
|
|
清理过期数据
|
|
|
|
Args:
|
|
days: 保留天数，如果为None则使用配置默认值
|
|
|
|
Returns:
|
|
int: 清理的记录数
|
|
"""
|
|
from ..core.config_manager import get_config
|
|
|
|
config = get_config()
|
|
if days is None:
|
|
days = config.database.data_retention_days
|
|
|
|
logger.info(f"开始清理 {days} 天前的过期数据")
|
|
|
|
try:
|
|
deleted_count = self.db_manager.cleanup_expired_data(days)
|
|
|
|
if deleted_count > 0:
|
|
logger.info(f"成功清理 {deleted_count} 条过期数据")
|
|
else:
|
|
logger.info("没有找到需要清理的过期数据")
|
|
|
|
return deleted_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"清理过期数据失败: {str(e)}")
|
|
return 0
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""
|
|
获取存储统计信息
|
|
|
|
Returns:
|
|
Dict[str, Any]: 统计数据
|
|
"""
|
|
try:
|
|
stats = self.db_manager.get_statistics()
|
|
stats.update({
|
|
"storage_type": "postgresql",
|
|
"last_cleanup": datetime.now().isoformat()
|
|
})
|
|
return stats
|
|
except Exception as e:
|
|
logger.error(f"获取存储统计信息失败: {str(e)}")
|
|
return {
|
|
"storage_type": "postgresql",
|
|
"error": str(e),
|
|
"total_announcements": 0,
|
|
"last_cleanup": datetime.now().isoformat()
|
|
}
|
|
|
|
def search_announcements(self, keyword: Optional[str] = None,
|
|
source_code: Optional[str] = None,
|
|
start_date: Optional[datetime] = None,
|
|
end_date: Optional[datetime] = None,
|
|
limit: int = 50) -> List[Announcement]:
|
|
"""
|
|
搜索公告
|
|
|
|
Args:
|
|
keyword: 关键词
|
|
source_code: 来源代码
|
|
start_date: 开始日期
|
|
end_date: 结束日期
|
|
limit: 限制数量
|
|
|
|
Returns:
|
|
List[Announcement]: 搜索结果
|
|
"""
|
|
# 这里可以实现更复杂的搜索逻辑
|
|
# 目前使用现有的查询方法
|
|
try:
|
|
return self.db_manager.get_announcements(
|
|
source_code=source_code,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
limit=limit
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"搜索公告失败: {str(e)}")
|
|
return []
|
|
|
|
def is_healthy(self) -> bool:
|
|
"""
|
|
检查存储健康状态
|
|
|
|
Returns:
|
|
bool: 是否健康
|
|
"""
|
|
try:
|
|
# 尝试执行一个简单的查询
|
|
stats = self.get_statistics()
|
|
return "error" not in stats
|
|
except Exception as e:
|
|
logger.error(f"存储健康检查失败: {str(e)}")
|
|
return False
|
|
|
|
def optimize_storage(self):
|
|
"""优化存储性能"""
|
|
# 这里可以添加数据库优化逻辑，如重建索引、清理碎片等
|
|
logger.info("开始优化存储性能")
|
|
|
|
try:
|
|
# 执行一些基本的优化操作
|
|
# 注意：实际的优化命令取决于PostgreSQL版本和配置
|
|
|
|
# 这里可以添加具体的优化SQL
|
|
# 例如：VACUUM, REINDEX等
|
|
|
|
logger.info("存储优化完成")
|
|
except Exception as e:
|
|
logger.error(f"存储优化失败: {str(e)}")
|
|
|
|
def backup_data(self, backup_path: Optional[str] = None) -> bool:
|
|
"""
|
|
备份数据
|
|
|
|
Args:
|
|
backup_path: 备份文件路径
|
|
|
|
Returns:
|
|
bool: 备份是否成功
|
|
"""
|
|
# 这里可以实现数据备份逻辑
|
|
# 可以使用pg_dump或其他备份工具
|
|
|
|
logger.info("开始备份数据")
|
|
|
|
try:
|
|
# 实现备份逻辑
|
|
# 注意：这需要系统权限来执行pg_dump
|
|
|
|
logger.info("数据备份完成")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"数据备份失败: {str(e)}")
|
|
return False
|
|
|
|
|
|
class StorageManager:
|
|
"""存储管理器"""
|
|
|
|
def __init__(self):
|
|
self.postgresql = PostgreSQLStorage()
|
|
self._current_storage = self.postgresql # 默认使用PostgreSQL
|
|
|
|
def save_announcements(self, announcements: List[Announcement]) -> int:
|
|
"""保存筛选后的公告"""
|
|
return self._current_storage.save_announcements(announcements)
|
|
|
|
def save_all_announcements_by_source(self, announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""按来源保存所有公告"""
|
|
return self._current_storage.save_all_announcements_by_source(announcements, max_per_source)
|
|
|
|
def save_auto_announcements(self, announcements: List[Announcement]) -> int:
|
|
"""保存定时搜索公告到专用表"""
|
|
return self._current_storage.save_auto_announcements(announcements)
|
|
|
|
def save_manual_announcements_by_source(self, announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""按来源保存手动搜索公告到专用表"""
|
|
return self._current_storage.save_manual_announcements_by_source(announcements, max_per_source)
|
|
|
|
def save_crawl_results(self, results: List[CrawlResult]) -> int:
|
|
"""保存爬取结果"""
|
|
return self._current_storage.save_crawl_results(results)
|
|
|
|
def get_recent_announcements(self, hours: int = 24, limit: int = 100) -> List[Announcement]:
|
|
"""获取最近公告"""
|
|
return self._current_storage.get_recent_announcements(hours, limit)
|
|
|
|
def cleanup_expired_data(self, days: Optional[int] = None) -> int:
|
|
"""清理过期数据"""
|
|
return self._current_storage.cleanup_expired_data(days)
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""获取统计信息"""
|
|
return self._current_storage.get_statistics()
|
|
|
|
def is_healthy(self) -> bool:
|
|
"""检查健康状态"""
|
|
return self._current_storage.is_healthy()
|
|
|
|
def optimize(self):
|
|
"""优化存储"""
|
|
self._current_storage.optimize_storage()
|
|
|
|
def backup(self, backup_path: Optional[str] = None) -> bool:
|
|
"""备份数据"""
|
|
return self._current_storage.backup_data(backup_path)
|
|
|
|
|
|
# 全局存储管理器实例
|
|
_storage_manager = None
|
|
_storage_lock = threading.Lock()
|
|
|
|
|
|
def get_storage_manager() -> StorageManager:
|
|
"""
|
|
获取存储管理器实例
|
|
|
|
Returns:
|
|
StorageManager: 存储管理器实例
|
|
"""
|
|
global _storage_manager
|
|
if _storage_manager is None:
|
|
with _storage_lock:
|
|
if _storage_manager is None:
|
|
_storage_manager = StorageManager()
|
|
return _storage_manager
|
|
|
|
|
|
def init_storage():
|
|
"""初始化存储"""
|
|
try:
|
|
init_database()
|
|
logger.info("存储初始化完成")
|
|
except Exception as e:
|
|
logger.error(f"存储初始化失败: {str(e)}")
|
|
raise
|
|
|
|
|
|
def save_announcements_to_storage(announcements: List[Announcement]) -> int:
|
|
"""
|
|
保存筛选后的公告到存储
|
|
|
|
Args:
|
|
announcements: 公告列表
|
|
|
|
Returns:
|
|
int: 保存成功的数量
|
|
"""
|
|
return get_storage_manager().save_announcements(announcements)
|
|
|
|
|
|
def save_all_announcements_by_source_to_storage(announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""
|
|
按来源保存所有公告到存储
|
|
|
|
Args:
|
|
announcements: 所有公告列表
|
|
max_per_source: 每个来源最大保留数量
|
|
|
|
Returns:
|
|
Dict[str, int]: 各来源保存的数量
|
|
"""
|
|
return get_storage_manager().save_all_announcements_by_source(announcements, max_per_source)
|
|
|
|
|
|
def save_auto_announcements_to_storage(announcements: List[Announcement]) -> int:
|
|
"""
|
|
保存定时搜索公告到专用表
|
|
|
|
Args:
|
|
announcements: 公告列表
|
|
|
|
Returns:
|
|
int: 保存成功的数量
|
|
"""
|
|
return get_storage_manager().save_auto_announcements(announcements)
|
|
|
|
|
|
def save_manual_announcements_by_source_to_storage(announcements: List[Announcement],
|
|
max_per_source: int = 100) -> Dict[str, int]:
|
|
"""
|
|
按来源保存手动搜索公告到专用表
|
|
|
|
Args:
|
|
announcements: 所有公告列表
|
|
max_per_source: 每个来源最大保留数量
|
|
|
|
Returns:
|
|
Dict[str, int]: 各来源保存的数量
|
|
"""
|
|
return get_storage_manager().save_manual_announcements_by_source(announcements, max_per_source)
|
|
|
|
|
|
def cleanup_storage(days: Optional[int] = None) -> int:
|
|
"""
|
|
清理存储中的过期数据
|
|
|
|
Args:
|
|
days: 保留天数
|
|
|
|
Returns:
|
|
int: 清理的记录数
|
|
"""
|
|
return get_storage_manager().cleanup_expired_data(days)
|