stock/stock-html/sync_fund_flow.py

#!/usr/bin/env python3
"""
资金流向数据每日采集脚本

功能：从5分钟K线数据自行计算资金流向，存入 stock_fund_flow_history 和 stock_fund_flow_today 表。
数据源：stock_kline_5min 表（自有数据，无需外部API）

算法：
  - 根据5分钟K线的 close vs open 判断买卖方向
  - 根据成交额(amount)分类：超大单(≥100万), 大单(20~100万), 中单(4~20万), 小单(<4万)
  - 主力 = 超大单 + 大单

用法:
  # 计算今日资金流向
  ./venv/bin/python sync_fund_flow.py

  # 补算历史（有5分钟K线但尚无资金流向的日期，最多30天）
  ./venv/bin/python sync_fund_flow.py --backfill

建议定时任务:
  10 15 * * 1-5 /opt/stock-app/venv/bin/python /opt/stock-app/sync_fund_flow.py >> /opt/stock-app/sync_fund_flow.log 2>&1
  30 15 * * 1-5 /opt/stock-app/venv/bin/python /opt/stock-app/sync_fund_flow.py --backfill >> /opt/stock-app/sync_fund_flow.log 2>&1
"""
import sys
import os
import time
import argparse
import fcntl
import atexit
from datetime import datetime, date
from collections import defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import psycopg2
from psycopg2.extras import execute_values
from config import Config

LOCK_FILE = '/tmp/sync_fund_flow.lock'
_lock_fd = None


def acquire_lock():
    """获取进程锁，防止多实例同时运行"""
    global _lock_fd
    _lock_fd = open(LOCK_FILE, 'w')
    try:
        fcntl.flock(_lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        _lock_fd.write(str(os.getpid()))
        _lock_fd.flush()
        atexit.register(release_lock)
        return True
    except IOError:
        try:
            with open(LOCK_FILE, 'r') as f:
                old_pid = f.read().strip()
            print(f"⚠️ 另一个实例正在运行 (PID: {old_pid})，退出", flush=True)
        except Exception:
            print(f"⚠️ 另一个实例正在运行，退出", flush=True)
        _lock_fd.close()
        _lock_fd = None
        return False


def release_lock():
    """释放进程锁"""
    global _lock_fd
    if _lock_fd:
        try:
            fcntl.flock(_lock_fd, fcntl.LOCK_UN)
            _lock_fd.close()
        except Exception:
            pass
        _lock_fd = None
    try:
        os.remove(LOCK_FILE)
    except Exception:
        pass


def get_db_conn():
    return psycopg2.connect(
        host=Config.DB_HOST, port=Config.DB_PORT,
        dbname=Config.DB_NAME, user=Config.DB_USER, password=Config.DB_PASSWORD,
    )


def calc_fund_flow_from_5min(conn, target_date):
    """
    从5分钟K线数据计算某日资金流向
    
    算法:
    1. 每根5分钟K线根据 close vs open 判断方向(买入/卖出)
    2. 根据成交额(amount)分类:
       - 超大单: amount >= 100万
       - 大单: 20万 <= amount < 100万
       - 中单: 4万 <= amount < 20万
       - 小单: amount < 4万
    3. 主力 = 超大单 + 大单
    """
    cur = conn.cursor()
    cur.execute("""
        SELECT code, open, close, volume, amount
        FROM stock_kline_5min
        WHERE dt::date = %s AND amount > 0
        ORDER BY code, dt
    """, (target_date,))
    rows = cur.fetchall()
    
    if not rows:
        return {}
    
    stock_flows = defaultdict(lambda: {
        'super_buy': 0, 'super_sell': 0,
        'big_buy': 0, 'big_sell': 0,
        'mid_buy': 0, 'mid_sell': 0,
        'small_buy': 0, 'small_sell': 0,
        'total_amount': 0
    })
    
    for code, open_p, close_p, volume, amount in rows:
        if not amount or float(amount) <= 0:
            continue
        
        amt = float(amount)
        sf = stock_flows[code]
        sf['total_amount'] += amt
        
        is_buy = float(close_p) >= float(open_p) if close_p and open_p else True
        
        if amt >= 1000000:      # 超大单 >= 100万
            cat = 'super'
        elif amt >= 200000:     # 大单 >= 20万
            cat = 'big'
        elif amt >= 40000:      # 中单 >= 4万
            cat = 'mid'
        else:                   # 小单
            cat = 'small'
        
        if is_buy:
            sf[f'{cat}_buy'] += amt
        else:
            sf[f'{cat}_sell'] += amt
    
    results = {}
    for code, sf in stock_flows.items():
        total = sf['total_amount']
        if total <= 0:
            continue
        
        super_net = sf['super_buy'] - sf['super_sell']
        big_net = sf['big_buy'] - sf['big_sell']
        mid_net = sf['mid_buy'] - sf['mid_sell']
        small_net = sf['small_buy'] - sf['small_sell']
        main_net = super_net + big_net
        
        results[code] = {
            'main_net_inflow': round(main_net, 2),
            'main_net_inflow_pct': round(main_net / total * 100, 4) if total > 0 else 0,
            'super_net_inflow': round(super_net, 2),
            'super_net_inflow_pct': round(super_net / total * 100, 4) if total > 0 else 0,
            'big_net_inflow': round(big_net, 2),
            'big_net_inflow_pct': round(big_net / total * 100, 4) if total > 0 else 0,
            'mid_net_inflow': round(mid_net, 2),
            'mid_net_inflow_pct': round(mid_net / total * 100, 4) if total > 0 else 0,
            'small_net_inflow': round(small_net, 2),
            'small_net_inflow_pct': round(small_net / total * 100, 4) if total > 0 else 0,
        }
    
    return results


def update_today_flow(conn):
    """更新今日资金流向到 stock_fund_flow_today"""
    today = date.today()
    flows = calc_fund_flow_from_5min(conn, today)
    
    if not flows:
        print(f"  今日({today})无5分钟K线数据，跳过", flush=True)
        return 0
    
    cur = conn.cursor()
    codes = list(flows.keys())
    cur.execute("""
        SELECT code, name, price, change_pct
        FROM stock_realtime_price WHERE code = ANY(%s)
    """, (codes,))
    price_map = {r[0]: {'name': r[1], 'price': float(r[2] or 0), 'change_pct': float(r[3] or 0)}
                 for r in cur.fetchall()}
    
    records = []
    for code, f in flows.items():
        info = price_map.get(code, {})
        records.append((
            code, info.get('name', ''),
            f['main_net_inflow'], f['main_net_inflow_pct'],
            f['super_net_inflow'], f['super_net_inflow_pct'],
            f['big_net_inflow'], f['big_net_inflow_pct'],
            f['mid_net_inflow'], f['mid_net_inflow_pct'],
            f['small_net_inflow'], f['small_net_inflow_pct'],
            info.get('price', 0), info.get('change_pct', 0),
        ))
    
    execute_values(cur, """
        INSERT INTO stock_fund_flow_today
        (code, name, main_net_inflow, main_net_inflow_pct,
         super_net_inflow, super_net_inflow_pct,
         big_net_inflow, big_net_inflow_pct,
         mid_net_inflow, mid_net_inflow_pct,
         small_net_inflow, small_net_inflow_pct,
         price, change_pct, updated_at)
        VALUES %s
        ON CONFLICT (code) DO UPDATE SET
            name = EXCLUDED.name,
            main_net_inflow = EXCLUDED.main_net_inflow,
            main_net_inflow_pct = EXCLUDED.main_net_inflow_pct,
            super_net_inflow = EXCLUDED.super_net_inflow,
            super_net_inflow_pct = EXCLUDED.super_net_inflow_pct,
            big_net_inflow = EXCLUDED.big_net_inflow,
            big_net_inflow_pct = EXCLUDED.big_net_inflow_pct,
            mid_net_inflow = EXCLUDED.mid_net_inflow,
            mid_net_inflow_pct = EXCLUDED.mid_net_inflow_pct,
            small_net_inflow = EXCLUDED.small_net_inflow,
            small_net_inflow_pct = EXCLUDED.small_net_inflow_pct,
            price = EXCLUDED.price,
            change_pct = EXCLUDED.change_pct,
            updated_at = NOW()
    """, records,
    template="(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW())")
    
    conn.commit()
    print(f"  ✅ 今日资金流向: {len(records)} 只股票", flush=True)
    return len(records)


def backfill_history(conn, max_days=30):
    """补算历史资金流向: 有5分钟K线但尚无 fund_flow_history 的日期"""
    cur = conn.cursor()
    cur.execute("""
        SELECT DISTINCT dt::date as d
        FROM stock_kline_5min
        WHERE dt::date NOT IN (
            SELECT DISTINCT trade_date FROM stock_fund_flow_history
        )
        AND dt::date < CURRENT_DATE
        ORDER BY d DESC
        LIMIT %s
    """, (max_days,))
    missing_dates = [row[0] for row in cur.fetchall()]
    
    if not missing_dates:
        print("  ✅ 历史资金流向已完整，无需补算", flush=True)
        return 0
    
    print(f"  需补算 {len(missing_dates)} 天的历史资金流向", flush=True)
    total_records = 0
    
    for d in missing_dates:
        flows = calc_fund_flow_from_5min(conn, d)
        if not flows:
            continue
        
        # 获取当天收盘价和涨跌幅（通过前一日收盘价计算）
        cur.execute("""
            SELECT k.code, k.close,
                   CASE WHEN prev.close > 0
                        THEN ROUND((k.close - prev.close) / prev.close * 100, 2)
                        ELSE 0 END AS change_pct
            FROM stock_kline_daily k
            LEFT JOIN LATERAL (
                SELECT close FROM stock_kline_daily
                WHERE code = k.code AND trade_date < k.trade_date
                ORDER BY trade_date DESC LIMIT 1
            ) prev ON true
            WHERE k.trade_date = %s AND k.code = ANY(%s)
        """, (d, list(flows.keys())))
        price_map = {r[0]: {'close': float(r[1] or 0), 'change_pct': float(r[2] or 0)}
                     for r in cur.fetchall()}
        
        records = []
        for code, f in flows.items():
            info = price_map.get(code, {})
            records.append((
                code, d,
                info.get('close', 0), info.get('change_pct', 0),
                f['main_net_inflow'], f['main_net_inflow_pct'],
                f['super_net_inflow'], f['super_net_inflow_pct'],
                f['big_net_inflow'], f['big_net_inflow_pct'],
                f['mid_net_inflow'], f['mid_net_inflow_pct'],
                f['small_net_inflow'], f['small_net_inflow_pct'],
            ))
        
        if records:
            execute_values(cur, """
                INSERT INTO stock_fund_flow_history
                (code, trade_date, close_price, change_pct,
                 main_net_inflow, main_net_inflow_pct,
                 super_net_inflow, super_net_inflow_pct,
                 big_net_inflow, big_net_inflow_pct,
                 mid_net_inflow, mid_net_inflow_pct,
                 small_net_inflow, small_net_inflow_pct,
                 updated_at)
                VALUES %s
                ON CONFLICT (code, trade_date) DO UPDATE SET
                    close_price = EXCLUDED.close_price,
                    change_pct = EXCLUDED.change_pct,
                    main_net_inflow = EXCLUDED.main_net_inflow,
                    main_net_inflow_pct = EXCLUDED.main_net_inflow_pct,
                    super_net_inflow = EXCLUDED.super_net_inflow,
                    super_net_inflow_pct = EXCLUDED.super_net_inflow_pct,
                    big_net_inflow = EXCLUDED.big_net_inflow,
                    big_net_inflow_pct = EXCLUDED.big_net_inflow_pct,
                    mid_net_inflow = EXCLUDED.mid_net_inflow,
                    mid_net_inflow_pct = EXCLUDED.mid_net_inflow_pct,
                    small_net_inflow = EXCLUDED.small_net_inflow,
                    small_net_inflow_pct = EXCLUDED.small_net_inflow_pct,
                    updated_at = NOW()
            """, records,
            template="(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW())")
            conn.commit()
            total_records += len(records)
            print(f"    {d}: {len(records)} 只股票", flush=True)
    
    print(f"  ✅ 历史补算完成: {total_records} 条记录, {len(missing_dates)} 天", flush=True)
    return total_records


def main():
    if not acquire_lock():
        sys.exit(1)
    
    parser = argparse.ArgumentParser(description='资金流向计算（从5分钟K线数据）')
    parser.add_argument('--backfill', action='store_true',
                        help='补算历史资金流向（有5分钟K线但尚无资金流向的日期）')
    parser.add_argument('--max-days', type=int, default=30,
                        help='历史补算最大天数（默认30）')
    args = parser.parse_args()
    
    today = date.today()
    print(f"{'='*60}", flush=True)
    print(f"💰 资金流向计算（来源: 5分钟K线数据）", flush=True)
    print(f"📅 日期: {today}", flush=True)
    print(f"{'='*60}", flush=True)
    
    conn = get_db_conn()
    start_time = time.time()
    
    try:
        # 始终计算今日
        print("\n📊 计算今日资金流向...", flush=True)
        today_count = update_today_flow(conn)
        
        # 如果指定了 --backfill，补算历史
        if args.backfill:
            print(f"\n📜 补算历史资金流向(最多{args.max_days}天)...", flush=True)
            hist_count = backfill_history(conn, args.max_days)
        else:
            hist_count = 0
        
        elapsed = time.time() - start_time
        
        # 显示数据库统计
        cur = conn.cursor()
        cur.execute("""
            SELECT count(*), count(DISTINCT code),
                   min(trade_date), max(trade_date),
                   count(DISTINCT trade_date)
            FROM stock_fund_flow_history
        """)
        cnt, codes, min_d, max_d, days = cur.fetchone()
        
        print(f"\n{'='*60}", flush=True)
        print(f"✅ 完成! 耗时: {elapsed:.1f}秒", flush=True)
        print(f"  今日: {today_count} 条 | 历史补算: {hist_count} 条", flush=True)
        print(f"\n💰 stock_fund_flow_history 统计:", flush=True)
        print(f"  总记录: {cnt:,} 条 | {codes:,} 只股票 | {days} 个交易日", flush=True)
        if min_d:
            print(f"  日期范围: {min_d} ~ {max_d}", flush=True)
        print(f"{'='*60}", flush=True)
        
    except Exception as e:
        print(f"❌ 错误: {e}", flush=True)
        import traceback
        traceback.print_exc()
    finally:
        conn.close()


if __name__ == '__main__':
    main()