ljwx-stock/backend/run_all_fetchers.sh

#!/bin/bash

################################################################################
# 一键抓取所有A股数据脚本
# 包括：公司信息、板块数据、财务报表、资金流、龙虎榜
################################################################################

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志目录
LOG_DIR="logs/fetcher_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$LOG_DIR"

# 打印带颜色的消息
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 打印分隔线
print_separator() {
    echo "================================================================================"
}

# 运行单个fetcher
run_fetcher() {
    local script=$1
    local description=$2
    local args=$3
    local log_file="$LOG_DIR/${script%.py}.log"

    print_separator
    log_info "开始执行: $description"
    log_info "脚本: $script"
    log_info "日志: $log_file"
    print_separator

    local start_time=$(date +%s)

    if python "$script" $args 2>&1 | tee "$log_file"; then
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        log_success "✅ $description 完成 (耗时: ${duration}秒)"
    else
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        log_error "❌ $description 失败 (耗时: ${duration}秒)"
        log_warning "查看日志: $log_file"

        # 询问是否继续
        read -p "是否继续执行下一个抓取器？(y/n) " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_warning "用户选择中止执行"
            exit 1
        fi
    fi

    echo ""
}

# 主程序
main() {
    local start_time=$(date +%s)

    print_separator
    echo -e "${GREEN}"
    echo "   ___   ____ __                __   ____        __           ______     __       __            "
    echo "  / _ | / __// /____  ______ __/ /  / __ \\___ _/ /____ _    / ____/__ _/ /_____/ /  ___ ______ "
    echo " / __ |/_\\ \\ / __/ _ \\/ __/ // / _ \\/ /_/ / _ \`/ __/ _ \`/  / /_/ -_) __/ __/ __/ _ \\/ -_) __(_-<"
    echo "/_/ |_/___//_/  \\___/\\__/\\_,_/_//_/\\____/\\_,_/\\__/\\_,_/  \\__/\\__/\\__/\\__/\\__/_//_/\\__/_/ /___/"
    echo -e "${NC}"
    print_separator
    log_info "开始时间: $(date '+%Y-%m-%d %H:%M:%S')"
    log_info "日志目录: $LOG_DIR"
    print_separator
    echo ""

    # 检查虚拟环境
    if [[ -z "$VIRTUAL_ENV" ]]; then
        log_warning "未检测到虚拟环境，尝试激活..."
        if [[ -f ".venv/bin/activate" ]]; then
            source .venv/bin/activate
            log_success "虚拟环境已激活"
        else
            log_error "未找到虚拟环境，请先运行: python3 -m venv .venv && source .venv/bin/activate"
            exit 1
        fi
    fi

    # 检查数据库连接
    log_info "检查数据库连接..."
    if python -c "from app.config import get_settings; import psycopg2; settings = get_settings(); conn = psycopg2.connect(settings.database_url); conn.close()" 2>/dev/null; then
        log_success "数据库连接正常"
    else
        log_error "数据库连接失败，请检查数据库配置"
        exit 1
    fi
    echo ""

    # 询问执行模式
    echo "请选择执行模式："
    echo "  1) 完整模式 - 抓取所有数据（预计7-9小时）"
    echo "  2) 测试模式 - 每个抓取器只抓取少量数据（预计5-10分钟）"
    echo "  3) 自定义 - 选择要执行的抓取器"
    read -p "请输入选项 (1/2/3): " mode
    echo ""

    case $mode in
        1)
            # 完整模式
            log_info "执行完整模式"
            run_fetcher "fetch_company_info_v2.py" "公司基础信息抓取 (5,295只股票)" ""
            run_fetcher "fetch_sector_data.py" "板块数据抓取 (行业+概念板块)" ""
            run_fetcher "fetch_financial_data.py" "财务报表抓取 (最近4季)" ""
            run_fetcher "fetch_money_flow.py" "资金流数据抓取 (700天)" "--mode incremental"
            run_fetcher "fetch_top_list.py" "龙虎榜数据抓取 (2年)" ""
            ;;
        2)
            # 测试模式
            log_info "执行测试模式"
            run_fetcher "fetch_company_info_v2.py" "公司基础信息抓取 (测试10只)" "--test"
            run_fetcher "fetch_sector_data.py" "板块数据抓取 (测试10个板块)" "--test"
            run_fetcher "fetch_financial_data.py" "财务报表抓取 (测试10只)" "--test"
            run_fetcher "fetch_money_flow.py" "资金流数据抓取 (测试10只)" "--test"
            run_fetcher "fetch_top_list.py" "龙虎榜数据抓取 (测试)" ""
            ;;
        3)
            # 自定义模式
            log_info "自定义模式"
            echo "请选择要执行的抓取器（多选用空格分隔，如: 1 3 5）："
            echo "  1) 公司基础信息"
            echo "  2) 板块数据"
            echo "  3) 财务报表"
            echo "  4) 资金流数据"
            echo "  5) 龙虎榜"
            read -p "请输入选项: " choices
            echo ""

            for choice in $choices; do
                case $choice in
                    1) run_fetcher "fetch_company_info_v2.py" "公司基础信息抓取" "" ;;
                    2) run_fetcher "fetch_sector_data.py" "板块数据抓取" "" ;;
                    3) run_fetcher "fetch_financial_data.py" "财务报表抓取" "" ;;
                    4) run_fetcher "fetch_money_flow.py" "资金流数据抓取" "--mode incremental" ;;
                    5) run_fetcher "fetch_top_list.py" "龙虎榜数据抓取" "" ;;
                    *) log_warning "无效选项: $choice" ;;
                esac
            done
            ;;
        *)
            log_error "无效的模式选择"
            exit 1
            ;;
    esac

    # 执行完成统计
    local end_time=$(date +%s)
    local total_duration=$((end_time - start_time))
    local hours=$((total_duration / 3600))
    local minutes=$(((total_duration % 3600) / 60))
    local seconds=$((total_duration % 60))

    print_separator
    log_success "🎉 所有抓取任务已完成！"
    log_info "结束时间: $(date '+%Y-%m-%d %H:%M:%S')"
    log_info "总耗时: ${hours}小时 ${minutes}分钟 ${seconds}秒"
    log_info "日志目录: $LOG_DIR"
    print_separator

    # 数据验证
    echo ""
    log_info "正在验证数据..."
    python - <<EOF
import psycopg2
from app.config import get_settings

settings = get_settings()
conn = psycopg2.connect(settings.database_url)
cursor = conn.cursor()

print("\n数据统计:")
print("=" * 80)

tables = [
    ('t_company_info', '公司信息'),
    ('t_sector_info', '板块信息'),
    ('t_sector_constituents', '板块成分股'),
    ('t_income_statement', '利润表'),
    ('t_balance_sheet', '资产负债表'),
    ('t_money_flow', '资金流'),
    ('t_top_list', '龙虎榜')
]

for table, desc in tables:
    try:
        cursor.execute(f"SELECT COUNT(*) FROM {table}")
        count = cursor.fetchone()[0]
        print(f"  {desc:20s} ({table:25s}): {count:>10,} 条记录")
    except Exception as e:
        print(f"  {desc:20s} ({table:25s}): 表不存在或查询失败")

print("=" * 80)
conn.close()
EOF

    echo ""
    log_success "✅ 数据抓取流程完成！"
}

# 捕获Ctrl+C
trap 'echo ""; log_warning "用户中断执行"; exit 130' INT

# 执行主程序
main