jiangcang_uo/uo_atlas_import.py

import os
import random
import time
from datetime import datetime, timedelta
import pymongo
from pymongo.errors import PyMongoError, ServerSelectionTimeoutError, BulkWriteError
# import pandas as pd
from config import atlas_config, mongo_config, atlas_table, mongo_table_uo, uo_city_pairs_old, uo_city_pairs_new


def import_flight_range_status(atlas_db, mongo_db, city_pair, create_at_begin_stamp, create_at_end_stamp, 
                               limit=0, max_retries=3, base_sleep=1.0, out_table=None):
    """
    从atlas查询指定城市对、时间范围的航班状态数据，写入mongo集合。
    :param atlas_db: atlas数据库连接
    :param mongo_db: mongo数据库连接
    :param city_pair: 城市对（例如："BJSBKK"）
    :param create_at_begin_stamp: 查询开始时间戳（秒级）
    :param create_at_end_stamp: 查询结束时间戳（秒级）
    :param limit: 限制返回结果数量（0表示不限制）
    :param max_retries: 最大重试次数
    :param base_sleep: 基础重试间隔（秒）
    :param out_table: 目标mongo集合名称（默认：mongo_table_uo）
    """
    for attempt in range(1, max_retries + 1):
        try:
            print(f"🔁 第 {attempt}/{max_retries} 次尝试查询")
            src_collection = atlas_db[atlas_table]      # 源集合(atlas)
            out_collection = mongo_db[mongo_table_uo]   # 目标集合(mongo)
            # 聚合查询管道
            pipeline = [
                {
                    "$match": {
                        "citypair": city_pair,
                        "from_date": {"$ne": None},
                        "flight_weight": {"$regex": r"^UO.*(0;0|1;20)$"},  # 使用$regex进行正则匹配
                        "create_at": {"$gte": create_at_begin_stamp, "$lte": create_at_end_stamp}
                    }
                },
                {
                    "$addFields": {
                        "create_at_beijing": {
                            "$add": [
                                {"$toDate": {"$multiply": ["$create_at", 1000]}},
                                8 * 60 * 60 * 1000  # 加8小时，保持Date类型
                            ]
                        }
                    }
                },
                {
                    "$sort": {"from_date": 1, "create_at": 1}
                }
            ]
            
            if limit > 0:
                pipeline.append({"$limit": limit})

            # 执行查询
            t1 = time.time()
            results = list(src_collection.aggregate(pipeline))
            t2 = time.time()
            rt = round(t2 - t1, 3)
            print(f"查询用时: {rt} 秒")

            # 写入结果（不落原表 _id/source_id，用业务字段做去重）
            upserted = 0
            matched = 0
            write_failed = 0
            batch_size = 500
            ops = []

            def flush_ops(ops_to_flush):
                nonlocal upserted, matched, write_failed
                if not ops_to_flush:
                    return
                try:
                    bulk_res = out_collection.bulk_write(ops_to_flush, ordered=False)
                    upserted += bulk_res.upserted_count
                    matched += bulk_res.matched_count
                except BulkWriteError as e:
                    details = e.details or {}
                    write_failed += len(details.get("writeErrors", []))
                    upserted += details.get("nUpserted", 0)
                    matched += details.get("nMatched", 0)
                except (ServerSelectionTimeoutError, PyMongoError) as e:
                    write_failed += len(ops_to_flush)
                    print(f"⚠️ Mongo 批量写入失败 {city_pair}: {e}")
                    if isinstance(e, ServerSelectionTimeoutError):
                        raise

            for i, doc in enumerate(results, start=1):
                citypair = doc.get('citypair')
                citypair_new = citypair[:3] + '-' + citypair[3:]
                from_date = doc.get('from_date')
                from_date_new = datetime.strptime(from_date, '%Y%m%d').strftime('%Y-%m-%d')
                flight_weight = doc.get('flight_weight')
                flight_weight_split = flight_weight.split(';')
                flight_numbers_raw = flight_weight_split[0]
                flight_numbers_parts = [x.strip() for x in flight_numbers_raw.split(',') if x.strip()]
                flight_numbers = ','.join(flight_numbers_parts)
                baggage_weight = int(flight_weight_split[-1])
                deptime = doc.get('depTime', '')
                if deptime:
                    from_time = datetime.strptime(deptime, '%Y%m%d%H%M').strftime('%Y-%m-%d %H:%M:%S')
                else:
                    from_time = ''   # 3月16之前抓的数据没有起飞时间
                trip_type = doc.get('trip_type')
                cabin_raw = doc.get('cabin', '')
                cabin_parts = [x.strip() for x in cabin_raw.split(',') if x.strip()]
                cabins = ','.join(cabin_parts)
                ticket_amount = doc.get('ticket_amount')
                currency = doc.get('vendorCurrency', '')
                price_base = doc.get('price')
                price_tax = doc.get('tax')
                price_total = doc.get('total')
                create_at = doc.get('create_at')
                create_at_beijing = doc.get('create_at_beijing')
                create_time = create_at_beijing.strftime('%Y-%m-%d %H:%M:%S')

                new_doc = {
                    "citypair": citypair_new,
                    "from_date": from_date_new,
                    "flight_numbers": flight_numbers,
                    "from_time": from_time,
                    "trip_type": trip_type,
                    "cabins": cabins,
                    "baggage_weight": baggage_weight,
                    "ticket_amount": ticket_amount,
                    "currency": currency,
                    "price_base": price_base,
                    "price_tax": price_tax,
                    "price_total": price_total,
                    "create_at": create_at,
                    "create_time": create_time,
                }

                dedup_filter = {
                    "citypair": citypair_new,
                    "from_date": from_date_new,
                    "create_at": create_at,
                    "flight_numbers": flight_numbers,
                    "baggage_weight": baggage_weight,
                }

                # try:
                #     res = out_collection.update_one(dedup_filter, {"$set": new_doc}, upsert=True)
                # except (ServerSelectionTimeoutError, PyMongoError) as e:
                #     write_failed += 1
                #     print(f"⚠️ Mongo 写入失败 {city_pair} [{i}/{len(results)}]: {e}")
                #     if isinstance(e, ServerSelectionTimeoutError):
                #         raise
                #     continue
                
                # if res.upserted_id is not None:
                #     upserted += 1
                # else:
                #     matched += res.matched_count

                ops.append(pymongo.UpdateOne(dedup_filter, {"$set": new_doc}, upsert=True))

                if len(ops) >= batch_size:
                    flush_ops(ops)
                    ops = []
        
                if i % 500 == 0:
                    print(f"写入进度 {city_pair} [{i}/{len(results)}]: upserted={upserted}, matched={matched}, failed={write_failed}")
            
            flush_ops(ops)

            print(f"写入集合: {mongo_table_uo}, upserted={upserted}, matched={matched}, failed={write_failed}, total={len(results)}")
            return

        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 mongo_con_parse(config=None):
    # if config is None:
    #     config = mongo_atlas_config.copy()
    try:
        if config.get("URI", ""):
            motor_uri = config["URI"]
            client = pymongo.MongoClient(motor_uri, maxPoolSize=100)
            db = client[config['db']]
            print("motor_uri: ", motor_uri)
        else:
            client = pymongo.MongoClient(
                config['host'],
                config['port'],
                serverSelectionTimeoutMS=30000,  # 30秒
                connectTimeoutMS=30000,  # 30秒
                socketTimeoutMS=30000,  # 30秒,
                retryReads=True,    # 开启重试
                maxPoolSize=50
            )
            db = client[config['db']]
            if config.get('user'):
                db.authenticate(config['user'], config['pwd'])

        print(f"✅ MongoDB 连接对象创建成功")

    except Exception as e:
        print(f"❌ 创建 MongoDB 连接对象时发生错误: {e}")
        raise
    return client, db

def main_import_process(create_at_begin, create_at_end):
    # 先转 datetime，再减去 0 小时得到 UTC 时间戳
    begin_dt = datetime.strptime(create_at_begin, "%Y-%m-%d %H:%M:%S") - timedelta(hours=0)
    end_dt = datetime.strptime(create_at_end, "%Y-%m-%d %H:%M:%S") - timedelta(hours=0)

    create_at_begin_stamp = int(begin_dt.timestamp())
    create_at_end_stamp = int(end_dt.timestamp())

    print(f"create_at_begin: {create_at_begin}, timestamp: {create_at_begin_stamp}")
    print(f"create_at_end: {create_at_end}, timestamp: {create_at_end_stamp}")

    uo_city_pairs = uo_city_pairs_new.copy()

    # 调试分支
    # uo_city_pairs = uo_city_pairs[5:6]

    for idx, city_pair in enumerate(uo_city_pairs):
        atlas_client, atlas_db = mongo_con_parse(atlas_config)
        mongo_client, mongo_db = mongo_con_parse(mongo_config)
        
        print(f"开始处理航线 {idx+1}/{len(uo_city_pairs)}: {city_pair}")    
        import_flight_range_status(atlas_db, mongo_db, city_pair, create_at_begin_stamp, create_at_end_stamp)
        print(f"结束处理航线 {idx+1}/{len(uo_city_pairs)}: {city_pair}")

        atlas_client.close()
        mongo_client.close()
    pass
    print("整体结束")
    

if __name__ == "__main__":
    print(f"本次导入开始时间: {datetime.now()}")
    
    current_time = datetime.now()
    create_at_end = current_time.strftime("%Y-%m-%d %H:%M:%S")
    create_at_begin = (current_time - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S")

    # create_at_begin = "2026-04-21 00:00:00"
    # create_at_end = "2026-04-22 23:59:59"

    main_import_process(create_at_begin, create_at_end)
    
    print(f"本次导入结束时间: {datetime.now()}")
    print()