stock/stock-html/sync_kline.py

#!/usr/bin/env python3
"""
K线数据本地同步脚本
支持四种数据源（按优先级）：阿里云API → 新浪API → 麦蕊API → AKShare
北交所(8XX/9XX)股票专用新浪API
同步到本地PostgreSQL数据库

用法:
  python sync_kline.py              # 增量同步（默认，只同步最近5天）
  python sync_kline.py --full       # 全量同步（180天历史数据）
  python sync_kline.py --days 30    # 同步最近30天
"""

import sys
import os
import time
import argparse
import signal as sig_module
from datetime import datetime, date, timedelta
from concurrent.futures import ThreadPoolExecutor, as_completed

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

import psycopg2
from psycopg2.extras import execute_values

from config import Config
from services.stock_algorithms import fetch_kline_rows

# ============ 配置 ============
WORKERS = 10            # 并发线程数
BATCH_SAVE_SIZE = 50    # 每批保存到DB的股票数
FULL_SYNC_DAYS = 180    # 全量同步天数
INCREMENTAL_DAYS = 5    # 增量同步天数（多取几天防遗漏）

_shutdown = False


def signal_handler(signum, frame):
    global _shutdown
    print("\n⚠️ 收到中断信号，正在优雅退出...")
    _shutdown = True


sig_module.signal(sig_module.SIGINT, signal_handler)
sig_module.signal(sig_module.SIGTERM, signal_handler)


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 get_all_stock_codes(conn):
    """获取可交易的股票列表（排除退市、停牌等无效股票）"""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT code, name FROM stock_realtime_price
            WHERE volume > 0 AND price > 0
              AND name NOT LIKE '%%退%%'
              AND name NOT LIKE 'PT%%'
            ORDER BY code
        """)
        return cur.fetchall()


def fetch_kline_for_stock(code, days):
    """获取K线数据 — 委托给 services.stock_algorithms.fetch_kline_rows"""
    return fetch_kline_rows(code, days)


def save_kline_batch(conn, all_rows):
    """批量保存K线数据到数据库（UPSERT）"""
    if not all_rows:
        return 0

    with conn.cursor() as cur:
        execute_values(
            cur,
            """
            INSERT INTO stock_kline_daily (code, trade_date, open, high, low, close, volume, amount)
            VALUES %s
            ON CONFLICT (code, trade_date) DO UPDATE SET
                open = EXCLUDED.open,
                high = EXCLUDED.high,
                low = EXCLUDED.low,
                close = EXCLUDED.close,
                volume = EXCLUDED.volume,
                amount = EXCLUDED.amount,
                updated_at = CURRENT_TIMESTAMP
            """,
            all_rows,
            page_size=1000,
        )
    conn.commit()
    return len(all_rows)


def main():
    global _shutdown

    parser = argparse.ArgumentParser(description='K线数据本地同步')
    parser.add_argument('--full', action='store_true', help='全量同步（180天历史）')
    parser.add_argument('--days', type=int, default=None, help='同步天数')
    args = parser.parse_args()

    if args.full:
        sync_days = FULL_SYNC_DAYS
        mode = '全量同步'
    elif args.days:
        sync_days = args.days
        mode = f'自定义同步（{sync_days}天）'
    else:
        sync_days = INCREMENTAL_DAYS
        mode = '增量同步'

    print(f"{'='*60}", flush=True)
    print(f"📊 K线数据本地同步", flush=True)
    print(f"📅 日期: {date.today()}", flush=True)
    print(f"🔄 模式: {mode}（{sync_days}天）", flush=True)
    print(f"⚙️  并发数: {WORKERS}", flush=True)
    print(f"{'='*60}", flush=True)

    conn = get_db_conn()
    all_stocks = get_all_stock_codes(conn)
    total = len(all_stocks)
    print(f"📈 股票总数: {total}", flush=True)

    start_time = time.time()
    done_count = 0
    success_count = 0
    error_count = 0
    total_rows = 0
    save_buffer = []
    last_report_time = time.time()

    print(f"\n🚀 开始同步...", flush=True)
    print(f"-" * 60, flush=True)

    with ThreadPoolExecutor(max_workers=WORKERS) as executor:
        futures = {}
        for code, name in all_stocks:
            if _shutdown:
                break
            futures[executor.submit(fetch_kline_for_stock, code, sync_days)] = (code, name)

        for future in as_completed(futures):
            if _shutdown:
                print("⏹️ 用户中断，正在保存当前数据...", flush=True)
                break

            code, name = futures[future]
            done_count += 1

            try:
                rows = future.result()
            except Exception:
                rows = None

            if rows:
                save_buffer.extend(rows)
                success_count += 1
            else:
                error_count += 1

            # 攒够一批就保存
            if len(save_buffer) >= BATCH_SAVE_SIZE * 80:  # 约50股 × 80条/股
                saved = save_kline_batch(conn, save_buffer)
                total_rows += saved
                save_buffer = []

            # 每3秒报告进度
            now = time.time()
            if now - last_report_time >= 3:
                elapsed = now - start_time
                speed = done_count / elapsed if elapsed > 0 else 0
                remaining = (total - done_count) / speed if speed > 0 else 0
                pct = done_count / total * 100
                print(f"  [{pct:5.1f}%] {done_count}/{total} "
                      f"| 速度: {speed:.1f}只/秒 | 剩余: {remaining/60:.1f}分钟 "
                      f"| 成功: {success_count} | 失败: {error_count} "
                      f"| 已保存: {total_rows}条", flush=True)
                last_report_time = now

    # 保存剩余数据
    if save_buffer:
        saved = save_kline_batch(conn, save_buffer)
        total_rows += saved

    elapsed = time.time() - start_time
    speed = done_count / elapsed if elapsed > 0 else 0

    print(f"\n{'='*60}", flush=True)
    print(f"✅ 同步{'中断' if _shutdown else '完成'}！", flush=True)
    print(f"  处理: {done_count} 只 | 成功: {success_count} | 失败: {error_count}", flush=True)
    print(f"  保存K线: {total_rows} 条 | 耗时: {elapsed/60:.1f}分钟", flush=True)
    print(f"  平均速度: {speed:.1f} 只/秒", flush=True)

    # 显示数据库统计
    with conn.cursor() as cur:
        cur.execute("SELECT count(*), count(DISTINCT code), min(trade_date), max(trade_date) FROM stock_kline_daily")
        cnt, codes, min_date, max_date = cur.fetchone()
        print(f"\n📊 数据库K线统计:", flush=True)
        print(f"  总记录: {cnt:,} 条 | 覆盖股票: {codes} 只", flush=True)
        print(f"  日期范围: {min_date} ~ {max_date}", flush=True)
    print(f"{'='*60}", flush=True)

    conn.close()


if __name__ == '__main__':
    main()