#!/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, '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()