FinDataHub/Scripts/bin/newsflash.py

#!/usr/bin/env python3
"""FinDataHub 快讯 — 查询 news_flash 数据库（支持来源/关键词/时间范围筛选+分页）"""
import json
import os
import sys
import argparse
from datetime import datetime, timedelta
from pathlib import Path

from common import MYSQL_CFG, MYSQL_CMD, MYSQL_PWD, DB_NEWS as DB

# 兼容：优先用 pymysql，没有则退回 mysql 命令行
try:
    import pymysql
    HAS_PYMYSQL = True
except ImportError:
    HAS_PYMYSQL = False
    import subprocess

# 连接池
from db_pool import get_conn, put_conn

MYSQL_CFG_DICT = MYSQL_CFG


def run_sql(sql, params=None):
    if HAS_PYMYSQL:
        try:
            conn = get_conn()
            cur = conn.cursor(pymysql.cursors.DictCursor)
            cur.execute(sql, params or ())
            rows = cur.fetchall()
            cur.close()
            put_conn(conn)
            str_rows = []
            for row in rows:
                str_row = {}
                for k, v in row.items():
                    if v is None:
                        str_row[k] = ""
                    else:
                        str_row[k] = str(v)
                str_rows.append(str_row)
            return str_rows
        except Exception as e:
            print(f"[sql error] {e}", file=sys.stderr)
            return None
    else:
        safe_sql = sql
        if params:
            for p in params:
                if isinstance(p, str):
                    safe_p = "'" + p.replace("'", "''") + "'"
                else:
                    safe_p = str(p)
                safe_sql = safe_sql.replace("%s", safe_p, 1)
        cmd = MYSQL_CMD + [DB, "-e", safe_sql]
        env = os.environ.copy()
        env["MYSQL_PWD"] = MYSQL_PWD
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=10, env=env)
        if r.returncode != 0:
            return None
        return _tabulate(r.stdout)


def _tabulate(stdout):
    """parse mysql tab-separated output（仅 fallback）"""
    lines = [l.rstrip("\r\n") for l in stdout.strip().split("\n") if l.strip()]
    if not lines:
        return []
    headers = lines[0].split("\t")
    rows = []
    for line in lines[1:]:
        vals = line.split("\t")
        rows.append(dict(zip(headers, vals)))
    return rows


def main():
    parser = argparse.ArgumentParser(description="查询快讯")
    parser.add_argument("--source", default="", help="来源 code（cls/wallstreetcn/sina/36kr/...）")
    parser.add_argument("--keyword", default="", help="标题+正文关键词搜索")
    parser.add_argument("--hours", type=int, default=0, help="近 N 小时（0=不限）")
    parser.add_argument("--limit", type=int, default=50, help="每页条数")
    parser.add_argument("--offset", type=int, default=0, help="分页偏移")
    args = parser.parse_args()

    result = {}

    conditions = ["f.published_at > 0"]
    params = []

    if args.source:
        conditions.append("s.code = %s")
        params.append(args.source)
    if args.keyword:
        like_kw = f"%{args.keyword}%"
        conditions.append("(f.title LIKE %s OR f.content LIKE %s)")
        params.extend([like_kw, like_kw])
    if args.hours > 0:
        cutoff_ms = int((datetime.now() - timedelta(hours=args.hours)).timestamp() * 1000)
        conditions.append("f.published_at >= %s")
        params.append(cutoff_ms)

    where = " AND ".join(conditions)

    # 1. 快讯列表（分页）
    list_sql = f"""
    SELECT s.name AS source_name, s.code AS source_code, f.title, f.content,
           f.summary, f.importance, f.published_at,
           FROM_UNIXTIME(f.published_at/1000) AS created_at
    FROM news_flash f
    LEFT JOIN news_sources s ON f.source_id = s.id
    WHERE {where}
    ORDER BY f.published_at DESC LIMIT %s OFFSET %s
    """
    list_params = params + [args.limit, args.offset]
    r = run_sql(list_sql, list_params)
    result["latest"] = r if r else []

    # 2. 总数（用于分页）
    count_sql = f"""
    SELECT COUNT(*) as total
    FROM news_flash f
    LEFT JOIN news_sources s ON f.source_id = s.id
    WHERE {where}
    """
    r = run_sql(count_sql, params)
    if r:
        result["total"] = int(r[0]["total"]) if r[0]["total"] else 0
    else:
        result["total"] = 0

    # 3. 各来源统计（仅在无筛选时返回，减少开销）
    if not args.source and not args.keyword and not args.hours:
        sql = """
        SELECT s.name AS source_name, COUNT(*) as count,
               MAX(FROM_UNIXTIME(f.published_at/1000)) as latest
        FROM news_flash f
        LEFT JOIN news_sources s ON f.source_id = s.id
        WHERE f.published_at > 0
        GROUP BY s.name ORDER BY count DESC
        """
        r = run_sql(sql)
        if r:
            result["sources"] = r

        # 4. 最近7天分布
        sql = """
        SELECT DATE(FROM_UNIXTIME(f.published_at/1000)) as day, COUNT(*) as count
        FROM news_flash f
        WHERE f.published_at > 0
        GROUP BY day ORDER BY day DESC LIMIT 7
        """
        r = run_sql(sql)
        if r:
            result["daily"] = r

        # 5. 今日条数 + 总条数
        sql = """
        SELECT COUNT(*) as today FROM news_flash
        WHERE published_at > 0
        AND published_at >= UNIX_TIMESTAMP(CURDATE()) * 1000
        """
        r = run_sql(sql)
        if r:
            result["today_count"] = int(r[0]["today"]) if r[0]["today"] else 0

        sql = "SELECT COUNT(*) as total FROM news_flash WHERE published_at > 0"
        r = run_sql(sql)
        if r:
            result["total_count"] = int(r[0]["total"]) if r[0]["total"] else 0

        # 热度统计
        sql = """
        SELECT importance, COUNT(*) as count
        FROM news_flash
        WHERE published_at > 0 AND importance >= 0
        GROUP BY importance ORDER BY importance DESC
        """
        r = run_sql(sql)
        if r:
            result["importance_dist"] = {row["importance"]: int(row["count"]) for row in r}

        # 高热度 Top5
        sql = """
        SELECT s.name AS source_name, f.title, f.summary, f.importance, f.published_at
        FROM news_flash f LEFT JOIN news_sources s ON f.source_id = s.id
        WHERE f.published_at > 0 AND f.importance >= 3
        ORDER BY f.importance DESC, f.published_at DESC LIMIT 5
        """
        r = run_sql(sql)
        if r:
            result["hot_top5"] = r

        # 今日总结
        sql = "SELECT summary, stats, updated_at FROM daily_summary WHERE date = CURDATE()"
        r = run_sql(sql)
        if r:
            result["daily_summary"] = {
                "summary": r[0]["summary"],
                "updated_at": r[0]["updated_at"],
            }

    result["status"] = "ok"
    print(json.dumps(result, ensure_ascii=False, default=str), flush=True)


if __name__ == "__main__":
    main()