yu-zhi-ran/scripts/db_helper.py

#!/usr/bin/env python3
"""
数据库辅助模块：为自动化脚本提供统一的数据库访问
"""

import sys
import os
from pathlib import Path
from datetime import datetime, date, timedelta
from typing import Optional, Dict, List

# 加载 .env（在 scripts/ 目录下运行时需要）
_env_path = Path(__file__).parent.parent / 'platform' / 'backend' / '.env'
if _env_path.exists():
    from dotenv import load_dotenv
    load_dotenv(_env_path)

# 添加项目根和 backend 路径
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'platform' / 'backend'))

from app.database import SessionLocal, init_db
from app.models import Topic
from sqlalchemy.orm import Session

# 确保数据库 schema 最新（含 org_id 等迁移）
init_db()

def get_topic_by_id(topic_id: str, db: Optional[Session] = None) -> Optional[Dict]:
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        topic = db.query(Topic).filter(Topic.id == topic_id).first()
        if not topic:
            return None
        return topic_to_dict(topic)
    finally:
        if close_db:
            db.close()

def get_topics_by_status(status: str, db: Optional[Session] = None) -> List[Dict]:
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        topics = db.query(Topic).filter(Topic.status == status).order_by(Topic.created_at).all()
        return [topic_to_dict(t) for t in topics]
    finally:
        if close_db:
            db.close()

def get_next_topic(priority: Optional[str] = None, db: Optional[Session] = None, today_only: bool = False) -> Optional[Dict]:
    """获取下一个待处理的选题（状态为 pending/待处理）"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        # 兼容两种状态表示
        status_filter = ['pending', '待处理']
        query = db.query(Topic).filter(Topic.status.in_(status_filter))
        if priority:
            query = query.filter(Topic.priority == priority)
        if today_only:
            cutoff = datetime.now() - timedelta(hours=24)
            query = query.filter(Topic.created_at >= cutoff)
        topic = query.order_by(Topic.priority_score.desc().nullslast(), Topic.created_at.asc()).first()
        return topic_to_dict(topic) if topic else None
    finally:
        if close_db:
            db.close()

def update_topic_status(topic_id: str, status: str, compliance_score: Optional[int] = None, db: Optional[Session] = None, reviewed_at: Optional[datetime] = None) -> bool:
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        topic = db.query(Topic).filter(Topic.id == topic_id).first()
        if not topic:
            return False
        topic.status = status
        topic.updated_at = datetime.now()
        if compliance_score is not None:
            topic.compliance_score = compliance_score
        if reviewed_at is not None:
            topic.reviewed_at = reviewed_at
        if status == 'review' and topic.generated_at is None:
            topic.generated_at = datetime.now()
            topic.reviewed_at = topic.reviewed_at or datetime.now()
        if status == 'ready':
            topic.ready_at = datetime.now().date()
            topic.reviewed_at = topic.reviewed_at or datetime.now()
        if status == 'published':
            topic.published_at = datetime.now().date()
        db.commit()
        return True
    finally:
        if close_db:
            db.close()

def get_active_llm_config(db: Optional[Session] = None) -> Optional[Dict]:
    """获取启用的 LLM 配置，优先读取 system_configs 中的 review_llm_id

    返回格式: { "model", "temperature", "max_tokens", "system_prompt", "name" }
    无配置时返回 None（调用方应使用环境变量默认值）
    """
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import LLMConfig, SystemConfig
        # 先查系统配置中指定的 review_llm_id
        sc = db.query(SystemConfig).filter(SystemConfig.key == 'review_llm_id').first()
        target_id = None
        if sc and sc.value and sc.value.isdigit():
            target_id = int(sc.value)
        query = db.query(LLMConfig)
        if target_id:
            query = query.filter(LLMConfig.id == target_id)
        query = query.filter(LLMConfig.is_active == True)
        cfg = query.first()
        if not cfg:
            cfg = db.query(LLMConfig).filter(LLMConfig.is_active == True).first()
        if cfg:
            return {
                "id": cfg.id,
                "name": cfg.name,
                "model": cfg.model,
                "temperature": cfg.temperature,
                "max_tokens": cfg.max_tokens,
                "system_prompt": cfg.system_prompt,
                "user_prompt_template": cfg.user_prompt_template,
            }
        return None
    finally:
        if close_db:
            db.close()

def topic_to_dict(topic: Topic) -> Dict:
    return {
        'id': topic.id,
        'title': topic.title,
        'field': topic.field_name,
        'format': topic.format,
        'core_concept': topic.core_concept,
        'audience_pain': topic.audience_pain,
        'unique_angle': topic.unique_angle,
        'priority': topic.priority,
        'priority_score': topic.priority_score or 0,
        'total_score': topic.total_score,
        'status': topic.status,
        'cases': topic.cases or [],
        'source_file': topic.source_file,
        'created_at': topic.created_at.isoformat() if topic.created_at else None,
        'updated_at': topic.updated_at.isoformat() if topic.updated_at else None,
        'ready_at': topic.ready_at.isoformat() if topic.ready_at else None,
        'published_at': topic.published_at.isoformat() if topic.published_at else None,
        'compliance_score': topic.compliance_score,
        'platform_urls': topic.platform_urls or {},
        'lock_by': None,
        'lock_at': None,
    }

def export_topics_to_json(db: Optional[Session] = None) -> List[Dict]:
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        topics = db.query(Topic).order_by(Topic.created_at).all()
        return [topic_to_dict(t) for t in topics]
    finally:
        if close_db:
            db.close()

if __name__ == "__main__":
    topics = export_topics_to_json()
    print(f"Total topics: {len(topics)}")
    for t in topics[:5]:
        print(f"- {t['id']}: {t['title'][:50]} ({t['status']})")

def save_topics_to_db(topics_data: List[Dict]):
    """保存/更新选题列表到数据库"""
    db = SessionLocal()
    try:
        for t in topics_data:
            existing = db.query(Topic).filter(Topic.id == t['id']).first()
            if existing:
                # 更新字段
                for field in ['title', 'field_name', 'format', 'core_concept', 'audience_pain', 'unique_angle', 'priority', 'priority_score', 'total_score', 'status', 'cases', 'source_file', 'compliance_score', 'platform_urls']:
                    setattr(existing, field, t.get(field, getattr(existing, field)))
                if t.get('ready_at'):
                    try:
                        existing.ready_at = datetime.strptime(t['ready_at'], '%Y-%m-%d').date()
                    except:
                        pass
                if t.get('published_at'):
                    try:
                        existing.published_at = datetime.strptime(t['published_at'], '%Y-%m-%d').date()
                    except:
                        pass
                existing.updated_at = datetime.now()
            else:
                new_topic = Topic(
                    id=t['id'],
                    title=t['title'],
                    field_name=t.get('field_name') or t.get('field', '可持续生活系统'),
                    format=t.get('format'),
                    core_concept=t.get('core_concept'),
                    audience_pain=t.get('audience_pain'),
                    unique_angle=t.get('unique_angle'),
                    priority=t.get('priority', '中'),
                    priority_score=t.get('priority_score', 0),
                    total_score=t.get('total_score'),
                    status=t.get('status', 'pending'),
                    cases=t.get('cases', []),
                    source_file=t.get('source_file'),
                    ready_at=datetime.strptime(t['ready_at'], '%Y-%m-%d').date() if t.get('ready_at') else None,
                    published_at=datetime.strptime(t['published_at'], '%Y-%m-%d').date() if t.get('published_at') else None,
                    compliance_score=t.get('compliance_score', 100),
                    platform_urls=t.get('platform_urls', {}),
                    created_at=datetime.now(),
                    updated_at=datetime.now()
                )
                db.add(new_topic)
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()

def save_article(topic_id: str, platform: str, html_content: str, db: Optional[Session] = None) -> Dict:
    """保存/更新文章到 articles 表"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import Article
        article_id = f"{platform}_{topic_id}"
        existing = db.query(Article).filter(Article.id == article_id).first()
        now = datetime.now()
        if existing:
            existing.html_content = html_content
        else:
            article = Article(
                id=article_id,
                topic_id=topic_id,
                platform=platform,
                file_path=f"db:{article_id}",
                html_content=html_content,
                status="draft",
                compliance_score=None
            )
            db.add(article)
        db.commit()
        return {"id": article_id, "topic_id": topic_id, "platform": platform}
    except Exception:
        db.rollback()
        raise
    finally:
        if close_db:
            db.close()

def save_cover_to_article(topic_id: str, platform: str, cover_path: str, db: Optional[Session] = None):
    """保存封面图路径到 article.images 字段"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import Article
        article_id = f"{platform}_{topic_id}"
        article = db.query(Article).filter(Article.id == article_id).first()
        if article:
            images = article.images or {}
            images["cover"] = cover_path
            article.images = images
            db.commit()
    except Exception:
        db.rollback()
    finally:
        if close_db:
            db.close()

def get_article(topic_id: str, platform: str, db: Optional[Session] = None) -> Optional[Dict]:
    """从 articles 表获取文章 HTML"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import Article
        article_id = f"{platform}_{topic_id}"
        article = db.query(Article).filter(Article.id == article_id).first()
        if not article:
            return None
        return {
            "id": article.id,
            "topic_id": article.topic_id,
            "platform": article.platform,
            "html_content": article.html_content,
            "status": article.status,
            "compliance_score": article.compliance_score,
            "created_at": article.created_at.isoformat() if article.created_at else None,
        }
    finally:
        if close_db:
            db.close()

def get_articles_by_topic(topic_id: str, db: Optional[Session] = None) -> List[Dict]:
    """获取某选题在所有平台的文章"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import Article
        articles = db.query(Article).filter(Article.topic_id == topic_id).all()
        return [{
            "id": a.id,
            "topic_id": a.topic_id,
            "platform": a.platform,
            "html_content": a.html_content,
            "status": a.status,
            "compliance_score": a.compliance_score,
        } for a in articles]
    finally:
        if close_db:
            db.close()

def delete_articles_by_topic(topic_id: str, db: Optional[Session] = None):
    """删除某选题的所有文章"""
    close_db = False
    if db is None:
        db = SessionLocal()
        close_db = True
    try:
        from app.models import Article
        db.query(Article).filter(Article.topic_id == topic_id).delete()
        db.commit()
    finally:
        if close_db:
            db.close()