auv/scripts/export_hotspot_data.py

#!/usr/bin/env python3
"""热点股追踪数据导出/导入脚本

用法：
  导出8月20日数据：
    python scripts/export_hotspot_data.py export --date 2026-08-20
  
  导入数据：
    python scripts/export_hotspot_data.py import --file hotspot_2026-08-20.json
  
  查看数据库中有哪些日期的数据：
    python scripts/export_hotspot_data.py list
"""

import argparse
import json
import os
import sys
from datetime import datetime
from typing import Optional

# 添加项目路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))

from database import get_connection, dict_from_row


def export_data(date: str, output_file: Optional[str] = None) -> dict:
    """导出指定日期的热点股追踪数据"""
    conn = get_connection()
    try:
        # 1. 导出每日题材前20
        themes = conn.execute(
            "SELECT * FROM daily_top_themes WHERE trade_date = ? ORDER BY rank ASC",
            (date,),
        ).fetchall()
        
        # 2. 导出每日核心股票
        core_stocks = conn.execute(
            "SELECT * FROM daily_core_stocks WHERE trade_date = ? ORDER BY rank ASC",
            (date,),
        ).fetchall()
        
        # 3. 导出每日核心股票与题材的关联
        stock_themes = conn.execute(
            "SELECT * FROM daily_core_stock_themes WHERE trade_date = ?",
            (date,),
        ).fetchall()
        
        data = {
            "version": "1.0",
            "export_time": datetime.now().isoformat(),
            "trade_date": date,
            "themes": [dict_from_row(r) for r in themes],
            "core_stocks": [dict_from_row(r) for r in core_stocks],
            "stock_themes": [dict_from_row(r) for r in stock_themes],
        }
        
        # 统计信息
        data["stats"] = {
            "theme_count": len(data["themes"]),
            "core_stock_count": len(data["core_stocks"]),
            "stock_theme_count": len(data["stock_themes"]),
        }
        
        # 保存到文件
        if output_file is None:
            output_file = f"hotspot_{date}.json"
        
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        print(f"✅ 导出成功：{output_file}")
        print(f"   题材数量：{data['stats']['theme_count']}")
        print(f"   核心股票：{data['stats']['core_stock_count']}")
        print(f"   股票-题材关联：{data['stats']['stock_theme_count']}")
        
        return data
    finally:
        conn.close()


def import_data(input_file: str, dry_run: bool = False) -> bool:
    """导入热点股追踪数据"""
    if not os.path.exists(input_file):
        print(f"❌ 文件不存在：{input_file}")
        return False
    
    with open(input_file, "r", encoding="utf-8") as f:
        data = json.load(f)
    
    # 验证数据格式
    required_keys = ["version", "trade_date", "themes", "core_stocks", "stock_themes"]
    if not all(k in data for k in required_keys):
        print("❌ 数据格式错误：缺少必要字段")
        return False
    
    date = data["trade_date"]
    print(f"📅 导入日期：{date}")
    print(f"   题材数量：{len(data['themes'])}")
    print(f"   核心股票：{len(data['core_stocks'])}")
    print(f"   股票-题材关联：{len(data['stock_themes'])}")
    
    if dry_run:
        print("\n🔍 试运行模式，不写入数据库")
        return True
    
    # 检查是否已存在
    conn = get_connection()
    try:
        existing = conn.execute(
            "SELECT 1 FROM daily_core_stocks WHERE trade_date = ? LIMIT 1",
            (date,),
        ).fetchone()
        
        if existing:
            print(f"\n⚠️  {date} 的数据已存在，将跳过重复数据（使用 INSERT OR IGNORE）")
    finally:
        conn.close()
    
    # 写入数据库
    conn = get_connection()
    try:
        # 导入每日题材前20
        for item in data["themes"]:
            # 移除自动生成的id和created_at
            item.pop("id", None)
            item.pop("created_at", None)
            conn.execute(
                "INSERT OR IGNORE INTO daily_top_themes (trade_date, theme_code, theme_name, bf3, hot_rank, rank) VALUES (?,?,?,?,?,?)",
                (item["trade_date"], item["theme_code"], item["theme_name"], item["bf3"], item["hot_rank"], item["rank"]),
            )
        
        # 导入每日核心股票
        for item in data["core_stocks"]:
            # 移除自动生成的id和created_at
            item.pop("id", None)
            item.pop("created_at", None)
            conn.execute(
                "INSERT OR IGNORE INTO daily_core_stocks (trade_date, stock_code, stock_name, f3, cover_count, rank) VALUES (?,?,?,?,?,?)",
                (item["trade_date"], item["stock_code"], item["stock_name"], item["f3"], item["cover_count"], item["rank"]),
            )
        
        # 导入每日核心股票与题材的关联
        for item in data["stock_themes"]:
            # 移除自动生成的id和created_at
            item.pop("id", None)
            item.pop("created_at", None)
            conn.execute(
                "INSERT OR IGNORE INTO daily_core_stock_themes (trade_date, stock_code, theme_code, theme_name) VALUES (?,?,?,?)",
                (item["trade_date"], item["stock_code"], item["theme_code"], item["theme_name"]),
            )
        
        conn.commit()
        print(f"\n✅ 导入成功：{date}")
        return True
    except Exception as e:
        conn.rollback()
        print(f"\n❌ 导入失败：{e}")
        return False
    finally:
        conn.close()


def list_dates() -> list[str]:
    """列出数据库中所有有数据的日期"""
    conn = get_connection()
    try:
        rows = conn.execute(
            "SELECT DISTINCT trade_date FROM daily_core_stocks ORDER BY trade_date DESC"
        ).fetchall()
        return [row["trade_date"] for row in rows]
    finally:
        conn.close()


def main():
    parser = argparse.ArgumentParser(description="热点股追踪数据导出/导入工具")
    subparsers = parser.add_subparsers(dest="command", help="可用命令")
    
    # 导出命令
    export_parser = subparsers.add_parser("export", help="导出数据")
    export_parser.add_argument("--date", required=True, help="交易日期 (YYYY-MM-DD)")
    export_parser.add_argument("--output", "-o", help="输出文件路径")
    
    # 导入命令
    import_parser = subparsers.add_parser("import", help="导入数据")
    import_parser.add_argument("--file", "-f", required=True, help="输入文件路径")
    import_parser.add_argument("--dry-run", action="store_true", help="试运行，不写入数据库")
    
    # 列表命令
    subparsers.add_parser("list", help="列出所有有数据的日期")
    
    args = parser.parse_args()
    
    if args.command == "export":
        export_data(args.date, args.output)
    elif args.command == "import":
        import_data(args.file, args.dry_run)
    elif args.command == "list":
        dates = list_dates()
        if dates:
            print("📅 数据库中的日期：")
            for d in dates:
                print(f"   {d}")
        else:
            print("📭 数据库中暂无数据")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()