from pymongo import MongoClient, ReadPreference
import time

# 配置
PRIMARY_HOST = "192.168.124.10:37017"
SECONDARY_HOST = "192.168.124.11:37017"
REPLICA_SET_NAME = "rs0"
CONN_STR = f"mongodb://{PRIMARY_HOST},{SECONDARY_HOST}/?replicaSet={REPLICA_SET_NAME}"

def get_server_address_for_query(collection, query):
    """
    执行查询，并返回处理该查询的服务器地址。
    这是最直接的诊断方法。
    """
    # 使用 explain() 获取查询执行的详细信息，其中包含 'serverInfo'
    explain_result = collection.find(query).explain()
    # 关键：提取实际处理查询的服务器地址
    server_address = explain_result.get('serverInfo', {}).get('host')
    return server_address

def accurate_read_route_test():
    """精确测试读请求被路由到了哪个节点"""
    print("=== 精确路由诊断测试 ===")
    
    # 1. 创建两个客户端
    write_client = MongoClient(CONN_STR)  # 用于写，默认连主库
    read_client = MongoClient(
        CONN_STR,
        read_preference=ReadPreference.SECONDARY_PREFERRED,
        localThresholdMS=50
    )
    
    db_write = write_client['test_route']
    db_read = read_client['test_route']
    coll_write = db_write['diagnostic']
    coll_read = db_read['diagnostic']
    
    # 清理并插入一条种子数据
    coll_write.delete_many({})
    seed_id = coll_write.insert_one({"test": "seed"}).inserted_id
    print("✓ 种子数据已插入\n")
    
    # 2. 进行多次读请求，并记录每个请求实际到达的服务器
    route_counter = {}  # 用于统计 {“主机地址”: 次数}
    test_iterations = 50
    
    print(f"开始执行 {test_iterations} 次查询，并捕获路由目标：")
    for i in range(test_iterations):
        # 每次读之前，先更新数据以生成新的oplog条目
        coll_write.update_one({"_id": seed_id}, {"$inc": {"counter": 1}})
        
        # 执行一次查询，并获取其实际执行的服务器地址
        target_server = get_server_address_for_query(coll_read, {"_id": seed_id})
        
        # 记录统计
        route_counter[target_server] = route_counter.get(target_server, 0) + 1
        
        print(f"  第{i+1:2d}次查询 -> 由服务器处理: {target_server}")
        time.sleep(0.1)  # 短暂间隔
    
    # 3. 输出统计结果
    print("\n=== 路由统计结果 ===")
    total_requests = sum(route_counter.values())
    for server, count in route_counter.items():
        percentage = (count / total_requests) * 100
        node_type = "主库 (PRIMARY)" if server == PRIMARY_HOST else "从库 (SECONDARY)"
        print(f"  {server} ({node_type}): {count} 次 ({percentage:.1f}%)")
    
    # 4. 清理
    write_client.close()
    read_client.close()
    print("\n诊断完成。")

if __name__ == "__main__":
    accurate_read_route_test()