Doiiars
323 lines
13 KiB
Python
323 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
BroadTopicExtraction模块 - 数据库管理器
|
||
只负责新闻数据和话题分析的存储和查询
|
||
"""
|
||
|
||
import sys
|
||
import json
|
||
from datetime import datetime, date, timedelta
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional
|
||
from sqlalchemy import create_engine, text
|
||
from sqlalchemy.engine import Engine
|
||
from loguru import logger
|
||
|
||
# 添加项目根目录到路径
|
||
project_root = Path(__file__).parent.parent
|
||
sys.path.append(str(project_root))
|
||
|
||
try:
|
||
import config
|
||
except ImportError:
|
||
raise ImportError("无法导入config.py配置文件")
|
||
|
||
from config import settings
|
||
|
||
|
||
class DatabaseManager:
|
||
"""数据库管理器"""
|
||
|
||
def __init__(self):
|
||
"""初始化数据库管理器"""
|
||
self.engine: Engine = None
|
||
self.connect()
|
||
|
||
def connect(self):
|
||
"""连接数据库"""
|
||
try:
|
||
dialect = (settings.DB_DIALECT or "mysql").lower()
|
||
if dialect in ("postgresql", "postgres"):
|
||
url = f"postgresql+psycopg://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}"
|
||
else:
|
||
url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_NAME}?charset={settings.DB_CHARSET}"
|
||
self.engine = create_engine(url, future=True)
|
||
logger.info(f"成功连接到数据库: {settings.DB_NAME}")
|
||
except ModuleNotFoundError as e:
|
||
missing: str = str(e)
|
||
if "psycopg" in missing:
|
||
logger.error(
|
||
"数据库连接失败: 未安装PostgreSQL驱动 psycopg。请安装: psycopg[binary]。参考指令：uv pip install psycopg[binary]")
|
||
elif "pymysql" in missing:
|
||
logger.error("数据库连接失败: 未安装MySQL驱动 pymysql。请安装: pymysql。参考指令：uv pip install pymysql")
|
||
else:
|
||
logger.error(f"数据库连接失败(缺少驱动): {e}")
|
||
raise
|
||
except Exception as e:
|
||
logger.exception(f"数据库连接失败: {e}")
|
||
raise
|
||
|
||
def close(self):
|
||
"""关闭数据库连接"""
|
||
if self.engine:
|
||
self.engine.dispose()
|
||
logger.info("数据库连接已关闭")
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.close()
|
||
|
||
# ==================== 新闻数据操作 ====================
|
||
|
||
def save_daily_news(self, news_data: List[Dict], crawl_date: date = None) -> int:
|
||
"""
|
||
保存每日新闻数据，如果当天已有数据则覆盖
|
||
|
||
Args:
|
||
news_data: 新闻数据列表
|
||
crawl_date: 爬取日期，默认为今天
|
||
|
||
Returns:
|
||
保存的新闻数量
|
||
"""
|
||
if not crawl_date:
|
||
crawl_date = date.today()
|
||
|
||
current_timestamp = int(datetime.now().timestamp())
|
||
|
||
try:
|
||
saved_count = 0
|
||
# 先独立事务执行删除，防止后续插入失败导致无法清理
|
||
with self.engine.begin() as conn:
|
||
deleted = conn.execute(text("DELETE FROM daily_news WHERE crawl_date = :d"), {"d": crawl_date}).rowcount
|
||
if deleted and deleted > 0:
|
||
logger.info(f"覆盖模式：删除了当天已有的 {deleted} 条新闻记录")
|
||
|
||
# 逐条插入，单条失败不影响后续（每条独立事务）
|
||
for news_item in news_data:
|
||
try:
|
||
# news_item.get('id') 已经是完整的 news_id（格式：source_item_id）
|
||
# 为了支持同一条新闻在不同日期出现，将 crawl_date 加入到 news_id 中
|
||
base_news_id = news_item.get(
|
||
'id') or f"{news_item.get('source', 'unknown')}_rank_{news_item.get('rank', 0)}"
|
||
# 将日期格式化为字符串并加入到 news_id 中，确保全局唯一性
|
||
news_id = f"{base_news_id}_{crawl_date.strftime('%Y%m%d')}"
|
||
|
||
title_val = (news_item.get("title", "") or "")
|
||
if len(title_val) > 500:
|
||
title_val = title_val[:500]
|
||
with self.engine.begin() as conn:
|
||
conn.execute(
|
||
text(
|
||
"""
|
||
INSERT INTO daily_news (
|
||
news_id, source_platform, title, url, crawl_date,
|
||
rank_position, add_ts, last_modify_ts
|
||
) VALUES (:news_id, :source_platform, :title, :url, :crawl_date, :rank_position, :add_ts, :last_modify_ts)
|
||
"""
|
||
),
|
||
{
|
||
"news_id": news_id,
|
||
"source_platform": news_item.get("source", "unknown"),
|
||
"title": title_val,
|
||
"url": news_item.get("url", ""),
|
||
"crawl_date": crawl_date,
|
||
"rank_position": news_item.get("rank", None),
|
||
"add_ts": current_timestamp,
|
||
"last_modify_ts": current_timestamp,
|
||
},
|
||
)
|
||
saved_count += 1
|
||
except Exception as e:
|
||
logger.exception(f"保存单条新闻失败: {e}")
|
||
continue
|
||
logger.info(f"成功保存 {saved_count} 条新闻记录")
|
||
return saved_count
|
||
except Exception as e:
|
||
logger.exception(f"保存新闻数据失败: {e}")
|
||
return 0
|
||
|
||
def get_daily_news(self, crawl_date: date = None) -> List[Dict]:
|
||
"""
|
||
获取每日新闻数据
|
||
|
||
Args:
|
||
crawl_date: 爬取日期，默认为今天
|
||
|
||
Returns:
|
||
新闻列表
|
||
"""
|
||
if not crawl_date:
|
||
crawl_date = date.today()
|
||
|
||
query = (
|
||
"SELECT * FROM daily_news WHERE crawl_date = :d ORDER BY rank_position ASC"
|
||
)
|
||
with self.engine.connect() as conn:
|
||
result = conn.execute(text(query), {"d": crawl_date})
|
||
rows = result.mappings().all()
|
||
return rows
|
||
|
||
# ==================== 话题数据操作 ====================
|
||
|
||
def save_daily_topics(self, keywords: List[str], summary: str, extract_date: date = None) -> bool:
|
||
"""
|
||
保存每日话题分析
|
||
|
||
Args:
|
||
keywords: 话题关键词列表
|
||
summary: 新闻分析总结
|
||
extract_date: 提取日期，默认为今天
|
||
|
||
Returns:
|
||
是否保存成功
|
||
"""
|
||
if not extract_date:
|
||
extract_date = date.today()
|
||
|
||
current_timestamp = int(datetime.now().timestamp())
|
||
|
||
try:
|
||
keywords_json = json.dumps(keywords, ensure_ascii=False)
|
||
# 为了支持外键引用，topic_id 需要全局唯一，所以将日期加入到 topic_id 中
|
||
topic_id = f"summary_{extract_date.strftime('%Y%m%d')}"
|
||
|
||
with self.engine.begin() as conn:
|
||
check = conn.execute(
|
||
text("SELECT id FROM daily_topics WHERE extract_date = :d AND topic_id = :tid"),
|
||
{"d": extract_date, "tid": topic_id},
|
||
).first()
|
||
if check:
|
||
conn.execute(
|
||
text(
|
||
"UPDATE daily_topics SET keywords = :k, topic_description = :s, add_ts = :ts, last_modify_ts = :lmt, topic_name = :tn WHERE extract_date = :d AND topic_id = :tid"
|
||
),
|
||
{"k": keywords_json, "s": summary, "ts": current_timestamp, "lmt": current_timestamp,
|
||
"d": extract_date, "tid": topic_id, "tn": "每日新闻分析"},
|
||
)
|
||
logger.info(f"更新了 {extract_date} 的话题分析")
|
||
else:
|
||
conn.execute(
|
||
text(
|
||
"INSERT INTO daily_topics (extract_date, topic_id, topic_name, keywords, topic_description, add_ts, last_modify_ts) VALUES (:d, :tid, :tn, :k, :s, :ts, :lmt)"
|
||
),
|
||
{"d": extract_date, "tid": topic_id, "tn": "每日新闻分析", "k": keywords_json, "s": summary,
|
||
"ts": current_timestamp, "lmt": current_timestamp},
|
||
)
|
||
logger.info(f"保存了 {extract_date} 的话题分析")
|
||
return True
|
||
except Exception as e:
|
||
logger.exception(f"保存话题分析失败: {e}")
|
||
return False
|
||
|
||
def get_daily_topics(self, extract_date: date = None) -> Optional[Dict]:
|
||
"""
|
||
获取每日话题分析
|
||
|
||
Args:
|
||
extract_date: 提取日期，默认为今天
|
||
|
||
Returns:
|
||
话题分析数据，如果不存在返回None
|
||
"""
|
||
if not extract_date:
|
||
extract_date = date.today()
|
||
|
||
try:
|
||
with self.engine.connect() as conn:
|
||
result = conn.execute(text("SELECT * FROM daily_topics WHERE extract_date = :d"),
|
||
{"d": extract_date}).mappings().first()
|
||
if result:
|
||
result = dict(result) # 转为可变dict以支持item赋值
|
||
result["keywords"] = json.loads(result["keywords"]) if result.get("keywords") else []
|
||
return result
|
||
return None
|
||
except Exception as e:
|
||
logger.exception(f"获取话题分析失败: {e}")
|
||
return None
|
||
|
||
def get_recent_topics(self, days: int = 7) -> List[Dict]:
|
||
"""
|
||
获取最近几天的话题分析
|
||
|
||
Args:
|
||
days: 天数
|
||
|
||
Returns:
|
||
话题分析列表
|
||
"""
|
||
try:
|
||
start_date = date.today() - timedelta(days=days)
|
||
with self.engine.connect() as conn:
|
||
results = conn.execute(
|
||
text(
|
||
"""
|
||
SELECT * FROM daily_topics
|
||
WHERE extract_date >= :start_date
|
||
ORDER BY extract_date DESC
|
||
"""
|
||
),
|
||
{"start_date": start_date},
|
||
).mappings().all()
|
||
for r in results:
|
||
r["keywords"] = json.loads(r["keywords"]) if r.get("keywords") else []
|
||
return results
|
||
except Exception as e:
|
||
logger.exception(f"获取最近话题分析失败: {e}")
|
||
return []
|
||
|
||
# ==================== 统计查询 ====================
|
||
|
||
def get_summary_stats(self, days: int = 7) -> Dict:
|
||
"""获取统计摘要"""
|
||
try:
|
||
start_date = date.today() - timedelta(days=days)
|
||
with self.engine.connect() as conn:
|
||
news_stats = conn.execute(
|
||
text(
|
||
"""
|
||
SELECT crawl_date, COUNT(*) as news_count, COUNT(DISTINCT source_platform) as platforms_count
|
||
FROM daily_news
|
||
WHERE crawl_date >= :start_date
|
||
GROUP BY crawl_date
|
||
ORDER BY crawl_date DESC
|
||
"""
|
||
),
|
||
{"start_date": start_date},
|
||
).all()
|
||
topics_stats = conn.execute(
|
||
text(
|
||
"""
|
||
SELECT extract_date, keywords, CHAR_LENGTH(topic_description) as summary_length
|
||
FROM daily_topics
|
||
WHERE extract_date >= :start_date
|
||
ORDER BY extract_date DESC
|
||
"""
|
||
),
|
||
{"start_date": start_date},
|
||
).all()
|
||
return {"news_stats": news_stats, "topics_stats": topics_stats}
|
||
except Exception as e:
|
||
logger.exception(f"获取统计摘要失败: {e}")
|
||
return {"news_stats": [], "topics_stats": []}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 测试数据库管理器
|
||
with DatabaseManager() as db:
|
||
# 测试获取新闻
|
||
news = db.get_daily_news()
|
||
logger.info(f"今日新闻数量: {len(news)}")
|
||
|
||
# 测试获取话题
|
||
topics = db.get_daily_topics()
|
||
if topics:
|
||
logger.info(f"今日话题关键词: {topics['keywords']}")
|
||
else:
|
||
logger.info("今日暂无话题分析")
|
||
|
||
logger.info("简化数据库管理器测试完成！")
|