yu-zhi-ran/scripts/db_helper.py

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

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

# 添加项目根和 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
from app.models import Topic
from sqlalchemy.orm import Session

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) -> 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)
        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, db: Optional[Session] = 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 status in ['ready', 'published'] and topic.generated_at is None:
            topic.generated_at = datetime.now()
        db.commit()
        return True
    finally:
        if close_db:
            db.close()

def topic_to_dict(topic: Topic) -> Dict:
    return {
        'id': topic.id,
        'title': topic.title,
        'field': topic.field.name if topic.field else None,
        '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', '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=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()