jiangcang_uo/data_loader.py

import time
import random
from datetime import datetime, timedelta
import pymongo
from pymongo.errors import PyMongoError, ServerSelectionTimeoutError
import pandas as pd
from uo_atlas_import import mongo_con_parse
from config import mongo_config, mongo_table_uo, uo_city_pairs


def query_groups_of_city_pair(db, city_pair, table_name, min_days=10, max_retries=3, base_sleep=1.0):
    """根据city_pair查询航线, 筛选1个月内至少有10天起飞的航线"""
    print(f"{city_pair} 查找所有分组")
    date_begin = (datetime.today() - timedelta(days=30)).strftime("%Y-%m-%d")
    date_end = datetime.today().strftime("%Y-%m-%d")
    
    # 聚合查询管道
    pipeline = [
        {
            "$match": {
                "citypair": city_pair,
                "from_date": {
                    "$gte": date_begin,
                    "$lte": date_end
                }
            }
        },
        {
            "$group": {
                "_id": {
                    "flight_numbers": "$flight_numbers",
                    "from_date": "$from_date"
                }
            }
        },
        {
            "$group": {
                "_id": "$_id.flight_numbers",
                "days": {"$sum": 1},
                "details": {"$push": "$_id.from_date"}
            }
        },
        {
            "$match": {
                "days": {"$gte": min_days}
            }
        },
        {
            "$addFields": {
                "details": {"$sortArray": {"input": "$details", "sortBy": 1}}
            }
        },
        {
            "$sort": {"_id": 1}
        }
    ]
    for attempt in range(1, max_retries + 1):
        try:
            print(f"  第 {attempt}/{max_retries} 次尝试查询")

            # 执行聚合查询
            collection = db[table_name]
            results = list(collection.aggregate(pipeline))

            # 格式化结果，使字段名更清晰
            formatted_results = [
                {
                    "flight_numbers": r["_id"],
                    "days": r["days"],
                    "flight_dates": r["details"]
                }
                for r in results
            ]
            
            return formatted_results

        except (ServerSelectionTimeoutError, PyMongoError) as e:
            print(f"⚠️ Mongo 查询失败: {e}")
            if attempt == max_retries:
                print("❌ 达到最大重试次数，放弃")
                return []

            # 指数退避 + 随机抖动
            sleep_time = base_sleep * (2 ** (attempt - 1)) + random.random()
            print(f"⏳ {sleep_time:.2f}s 后重试...")
            time.sleep(sleep_time)


def query_flight_range_status(db, table_name, city_pair, flight_numbers, from_date_begin, from_date_end, 
                              limit=0, max_retries=3, base_sleep=1.0):
    for attempt in range(1, max_retries + 1):
        try:
            print(f"🔁 第 {attempt}/{max_retries} 次尝试查询")
            # 构建查询条件
            projection = {
                # "_id": 0   # 一般会关掉
                "citypair": 1,
                "flight_numbers": 1,
                "from_date": 1,
                "from_time": 1,
                "create_time": 1,
                "baggage_weight": 1,
                "cabins": 1,
                "ticket_amount": 1,
                "currency": 1,
                "price_total": 1
            }            
            pipeline = [
                {
                    "$match": {
                        "citypair": city_pair,
                        "flight_numbers": flight_numbers,
                        "baggage_weight": {"$in": [0, 20]},
                        "from_date": {
                            "$gte": from_date_begin,
                            "$lte": from_date_end
                        }
                    }
                },
                {
                    "$project": projection   # 就是这里
                },
                {
                    "$sort": {
                        "from_date": 1,
                        "baggage_weight": 1, 
                        "create_time": 1
                    }
                }
            ]
            # print(f"   查询条件: {pipeline}")
            # 执行查询
            collection = db[table_name]
            results = list(collection.aggregate(pipeline))
            print(f"✅ 查询成功，找到 {len(results)} 条记录")

            if results:
                df = pd.DataFrame(results)
                if '_id' in df.columns:
                    df = df.drop(columns=['_id'])
                
                if 'from_time' in df.columns and 'from_date' in df.columns:
                    from_time_raw = df['from_time']
                    from_time_str = from_time_raw.fillna('').astype(str).str.strip()
                    non_empty = from_time_str[from_time_str.ne('')]  # 找到原始 from_time 非空的记录

                    extracted_time = non_empty.str.extract(r'(\d{2}:\d{2}:\d{2})$')[0].dropna()
                    if not extracted_time.empty:
                        more_time = extracted_time.value_counts().idxmax()   # 按众数分配给其它行 构造from_time
                        missing_mask = from_time_raw.isna() | from_time_str.eq('')
                        if missing_mask.any():
                            df.loc[missing_mask, 'from_time'] = df.loc[missing_mask, 'from_date'].astype(str).str.strip() + ' ' + more_time
                    else:
                        # 无法得到起飞日期的抛弃
                        return pd.DataFrame()

                print(f"📊 已转换为 DataFrame，形状: {df.shape}")
                return df

            else:
                print("⚠️  查询结果为空")
                return pd.DataFrame()

        except (ServerSelectionTimeoutError, PyMongoError) as e:
            print(f"⚠️ Mongo 查询失败: {e}")
            if attempt == max_retries:
                print("❌ 达到最大重试次数，放弃")
                return pd.DataFrame()
            
            # 指数退避 + 随机抖动
            sleep_time = base_sleep * (2 ** (attempt - 1)) + random.random()
            print(f"⏳ {sleep_time:.2f}s 后重试...")
            time.sleep(sleep_time)


def fill_hourly_create_time(df):
    """补齐成小时粒度数据"""
    pass

def process_flight_numbers(args):
    process_id, db_config, city_pair, flight_numbers, from_date_begin, from_date_end, is_train, plot_flag, output_dir = args
    print(f"[进程{process_id}] 开始处理航班号: {flight_numbers}")
    
    # 为每个进程创建独立的数据库连接
    try:
        client, db = mongo_con_parse(db_config)
        print(f"[进程{process_id}] ✅ 数据库连接创建成功")
    except Exception as e:
        print(f"[进程{process_id}] ❌ 数据库连接创建失败: {e}")
        return pd.DataFrame()
    
    try:
        # 查询
        df_1 = query_flight_range_status(db, mongo_table_uo, city_pair, flight_numbers, from_date_begin, from_date_end)
        
        df_f1 = fill_hourly_create_time(df_1)

        
    except Exception as e:
        print(f"[进程{process_id}] ❌ 处理航班号:{flight_numbers} 时发生异常: {e}")
        return pd.DataFrame()
    finally:
        # 确保关闭数据库连接
        try:
            client.close()
            print(f"[进程{process_id}] ✅ 数据库连接已关闭")
        except:
            pass


def load_data(db_config, city_pair, from_date_begin, from_date_end, is_train=True, plot_flag=False, output_dir='.', 
              use_multiprocess=False, max_workers=None):

    print(f"开始处理航线: {city_pair}")
    main_client, main_db = mongo_con_parse(db_config)
    all_groups = query_groups_of_city_pair(main_db, city_pair, mongo_table_uo)
    main_client.close()
    
    all_groups_len = len(all_groups)
    print(f"该航线共有{all_groups_len}组航班号")

    print("使用单进程处理")
    process_id = 0

    for each_group in all_groups:
        flight_numbers = each_group.get("flight_numbers", "未知")
        args = (process_id, db_config, city_pair, flight_numbers, from_date_begin, from_date_end, is_train, plot_flag, output_dir)
        try:
            df_mid = process_flight_numbers(args)
            pass

        except Exception as e:
            print(f"❌ 航班号:{flight_numbers} 处理异常: {e}")



if __name__ == "__main__":
    
    from_date_begin = "2026-03-17"
    from_date_end = "2026-04-01"

    uo_city_pair_list = [f"{pair[:3]}-{pair[3:]}" for pair in uo_city_pairs]

    for idx, uo_city_pair in enumerate(uo_city_pair_list, start=1):
        # 使用默认配置
        # client, db = mongo_con_parse()
        print(f"第 {idx} 组 :", uo_city_pair)

        start_time = time.time()
        load_data(mongo_config, uo_city_pair, from_date_begin, from_date_end)
        end_time = time.time()
        run_time = round(end_time - start_time, 3)
        print(f"用时: {run_time} 秒")