data_analysis/fetch_daily.py

#!/usr/bin/env python
# coding=utf-8
"""
按天增量获取数据 - 通用版本
支持并发获取，自动跳过已有数据

用法:
    python fetch_daily.py tasks/xxx/query.sql                    # 获取最近7天
    python fetch_daily.py tasks/xxx/query.sql --days 30          # 获取最近30天
    python fetch_daily.py tasks/xxx/query.sql --start 20260101 --end 20260107
    python fetch_daily.py tasks/xxx/query.sql --date 20260105    # 单天
    python fetch_daily.py tasks/xxx/query.sql --force            # 强制重新获取
    python fetch_daily.py tasks/xxx/query.sql --workers 10       # 设置天级并发数
    python fetch_daily.py tasks/xxx/query.sql --parallel 50      # 单天多线程下载（默认50，大数据量推荐）
    python fetch_daily.py tasks/xxx/query.sql --parallel 0       # 关闭多线程，使用单线程下载
"""
import argparse
import sys
from datetime import datetime, timedelta
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

sys.path.insert(0, str(Path(__file__).parent / "lib"))
from odps_module import ODPSClient

# 线程安全的计数器
counter_lock = threading.Lock()
success_count = 0
fail_count = 0


def get_existing_dates(daily_dir):
    """获取已下载的日期列表"""
    existing = set()
    if not daily_dir.exists():
        return existing
    for f in daily_dir.glob("*.csv"):
        try:
            dt = f.stem
            if len(dt) == 8 and dt.isdigit():
                existing.add(dt)
        except:
            pass
    return existing


def merge_csv_files(daily_dir, output_file=None):
    """合并目录下所有日期 CSV 文件，只保留一个表头"""
    csv_files = sorted(daily_dir.glob("*.csv"))
    if not csv_files:
        print("没有找到 CSV 文件")
        return None

    if output_file is None:
        output_file = daily_dir.parent / f"{daily_dir.name}_merged.csv"

    with open(output_file, "w", encoding="utf-8") as out:
        header_written = False
        total_rows = 0

        for csv_file in csv_files:
            with open(csv_file, "r", encoding="utf-8") as f:
                lines = f.readlines()
                if not lines:
                    continue

                if not header_written:
                    out.write(lines[0])
                    header_written = True

                for line in lines[1:]:
                    out.write(line)
                    total_rows += 1

    print(f"合并完成: {len(csv_files)} 个文件, {total_rows} 行数据")
    print(f"输出文件: {output_file}")
    return output_file


def get_date_range(start_str, end_str):
    """生成日期范围列表"""
    start = datetime.strptime(start_str, "%Y%m%d")
    end = datetime.strptime(end_str, "%Y%m%d")
    dates = []
    current = start
    while current <= end:
        dates.append(current.strftime("%Y%m%d"))
        current += timedelta(days=1)
    return dates


def fetch_single_day(dt, sql_template, daily_dir, parallel_threads=0):
    """获取单天数据"""
    global success_count, fail_count

    try:
        client = ODPSClient()
        sql = sql_template.replace("${dt}", dt)
        output_file = daily_dir / f"{dt}.csv"

        # 下载到文件
        if parallel_threads > 0:
            # 多线程并行下载（适合大数据量）
            client.execute_sql_result_save_file_parallel(sql, str(output_file), workers=parallel_threads)
        else:
            # 单线程下载
            client.execute_sql_result_save_file(sql, str(output_file))

        # 检查结果
        if output_file.exists():
            row_count = sum(1 for _ in open(output_file)) - 1  # 减去表头
            with counter_lock:
                success_count += 1
            if row_count > 0:
                return (dt, "success", row_count)
            else:
                return (dt, "empty", 0)
        else:
            with counter_lock:
                fail_count += 1
            return (dt, "fail", 0)

    except Exception as e:
        with counter_lock:
            fail_count += 1
        return (dt, "error", str(e))


def main():
    global success_count, fail_count

    parser = argparse.ArgumentParser(description="按天增量获取数据")
    parser.add_argument("sql_file", type=str, help="SQL文件路径")
    parser.add_argument("--days", type=int, default=7, help="获取最近N天 (默认7)")
    parser.add_argument("--start", type=str, help="开始日期 YYYYMMDD")
    parser.add_argument("--end", type=str, help="结束日期 YYYYMMDD")
    parser.add_argument("--date", type=str, help="单天日期 YYYYMMDD")
    parser.add_argument("--force", action="store_true", help="强制重新获取")
    parser.add_argument("--workers", type=int, default=5, help="天级并发数 (默认5)")
    parser.add_argument("--parallel", type=int, default=50, help="单天多线程下载 (默认50, 大数据量推荐)")
    parser.add_argument("--merge", action="store_true", help="合并所有日期数据到一个文件")
    args = parser.parse_args()

    # 解析 SQL 文件路径
    sql_file = Path(args.sql_file).resolve()
    if not sql_file.exists():
        print(f"错误: 找不到 {sql_file}")
        return

    # 输出目录：SQL 同目录下的 output/SQL文件名/
    output_dir = sql_file.parent / "output"
    daily_dir = output_dir / sql_file.stem
    daily_dir.mkdir(parents=True, exist_ok=True)

    print(f"SQL文件: {sql_file}")
    print(f"数据目录: {daily_dir}")

    # 仅合并模式：不获取数据，直接合并已有文件
    if args.merge:
        existing_dates = get_existing_dates(daily_dir)
        print(f"已有数据: {len(existing_dates)}天")
        if existing_dates:
            merge_csv_files(daily_dir)
        else:
            print("没有可合并的数据")
        return

    # 确定日期范围
    if args.date:
        target_dates = [args.date]
    elif args.start and args.end:
        target_dates = get_date_range(args.start, args.end)
    else:
        today = datetime.now()
        end_date = (today - timedelta(days=1)).strftime("%Y%m%d")
        start_date = (today - timedelta(days=args.days)).strftime("%Y%m%d")
        target_dates = get_date_range(start_date, end_date)

    print(f"目标日期: {target_dates[0]} ~ {target_dates[-1]} ({len(target_dates)}天)")

    # 检查已有数据
    existing_dates = get_existing_dates(daily_dir)
    print(f"已有数据: {len(existing_dates)}天")

    # 确定需要获取的日期
    if args.force:
        missing_dates = target_dates
        print(f"强制模式: 重新获取所有 {len(missing_dates)} 天")
    else:
        missing_dates = [d for d in target_dates if d not in existing_dates]
        print(f"需要获取: {len(missing_dates)}天")

    if not missing_dates:
        print("没有需要获取的数据，退出")
        return

    # 读取 SQL 模板
    sql_template = sql_file.read_text(encoding="utf-8")

    # 检测 SQL 中是否包含 ${dt} 变量
    has_dt_var = "${dt}" in sql_template

    # 重置计数器
    success_count = 0
    fail_count = 0

    # 如果 SQL 中没有 ${dt}，只需执行一次
    if not has_dt_var:
        print("\n检测到 SQL 中不含 ${dt} 变量，只执行一次...")
        target_dates = ["20000101"]  # 用虚拟日期
        missing_dates = target_dates
        output_file = output_dir / f"{sql_file.stem}.csv"
        output_file.parent.mkdir(parents=True, exist_ok=True)

        try:
            client = ODPSClient()
            if args.parallel > 0:
                client.execute_sql_result_save_file_parallel(sql_template, str(output_file), workers=args.parallel)
            else:
                client.execute_sql_result_save_file(sql_template, str(output_file))
            print(f"数据目录: {output_file}")
        except Exception as e:
            print(f"✗ 执行失败: {e}")
        return

    # 并发获取
    print(f"目标日期: {target_dates[0]} ~ {target_dates[-1]} ({len(target_dates)}天)")
    workers = min(args.workers, len(missing_dates))
    if args.parallel > 0:
        print(f"\n开始获取 (天级并发: {workers}, 单天多线程: {args.parallel})...")
    else:
        print(f"\n开始获取 (并发数: {workers})...")

    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(fetch_single_day, dt, sql_template, daily_dir, args.parallel): dt
            for dt in missing_dates
        }

        completed = 0
        for future in as_completed(futures):
            completed += 1
            dt, status, info = future.result()

            if status == "success":
                print(f"  [{completed}/{len(missing_dates)}] ✓ {dt}: {info} 行")
            elif status == "empty":
                print(f"  [{completed}/{len(missing_dates)}] ⚠ {dt}: 无数据")
            elif status == "error":
                print(f"  [{completed}/{len(missing_dates)}] ✗ {dt}: {info}")
            else:
                print(f"  [{completed}/{len(missing_dates)}] ✗ {dt}: 失败")

    print(f"\n完成! 成功: {success_count}, 失败: {fail_count}")
    print(f"数据目录: {daily_dir}")


if __name__ == "__main__":
    main()