请求优化OKTA用户-应用关联查询Python脚本，提升执行效率

问题背景 #

当前实现Okta全应用关联用户查询的Python脚本执行耗时近90分钟，平台有170个应用、约1100个活跃用户，已解决分页和限流问题，但效率极低。原脚本如下：

import requests
import csv
import os
import time

def get_all_users(domain, okta_api_token):
    headers = {
        'Authorization': f'SSWS {okta_api_token}',
        'Accept': 'application/json',
        'Content-Type': 'application/json'
    }
    
    okta_url = f'https://{domain}.okta.com'
    url = f'{okta_url}/api/v1/users'
    
    all_users = []
    while url:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        users = response.json()
        all_users += users
        
        links = response.links
        url = links.get('next', {}).get('url')
    
    return all_users

def get_user_apps(domain, okta_api_token, user_id, app_links):
    if user_id in app_links:
        return app_links[user_id]
    
    headers = {
        'Authorization': f'SSWS {okta_api_token}',
        'Accept': 'application/json',
        'Content-Type': 'application/json'
    }
    
    okta_url = f'https://{domain}.okta.com'
    url = f'{okta_url}/api/v1/users/{user_id}/appLinks'
    
    apps = []
    while url:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        app_links_data = response.json()
        apps += app_links_data
        
        links = response.links
        url = links.get('next', {}).get('url')
    
    app_names = [app['label'] for app in apps]
    app_links[user_id] = ', '.join(app_names)
    return app_links[user_id]

def extract_user_data(domain, okta_api_token, users, app_links):
    user_data = []
    for user in users:
        user_id = user.get('id', 'N/A')
        user_name = user.get('profile', {}).get('login', 'N/A')
        email = user.get('profile', {}).get('email', 'N/A')
        first_name = user.get('profile', {}).get('firstName', 'N/A')
        last_name = user.get('profile', {}).get('lastName', 'N/A')
        full_name = f"{first_name} {last_name}"
        
        apps = get_user_apps(domain, okta_api_token, user_id, app_links)
        user_data.append({'id': user_id, 'username': user_name, 'email': email, 'full_name': full_name, 'apps': apps})
    return user_data

def extract_app_user_data(domain, okta_api_token, users):
    app_user_data = {}
    app_links = {}
    for user in users:
        user_id = user.get('id', 'N/A')
        apps = get_user_apps(domain, okta_api_token, user_id, app_links)
        for app in apps.split(', '):
            if app not in app_user_data:
                app_user_data[app] = []
            app_user_data[app].append(user.get('profile', {}).get('login', 'N/A'))
    return app_user_data

# Read domain and okta_api_token from environment variables
domain = os.environ['OKTA_DOMAIN']
okta_api_token = os.environ['OKTA_API_TOKEN']

# Start the timer
start_time = time.time()

# Get all users and extract the desired data
users = get_all_users(domain, okta_api_token)
app_user_data = extract_app_user_data(domain, okta_api_token, users)
app_links = {}
extracted_users = extract_user_data(domain, okta_api_token, users, app_links)

# Save extracted app user data to a numbered CSV file
with open('app_user_data.csv', 'w', newline='') as csvfile:
    fieldnames = ['app', 'users']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()
    
    for app, users in app_user_data.items():
        writer.writerow({'app': app, 'users': ', '.join(users)})

# Save extracted user data to a numbered CSV file
with open('user_data.csv', 'w', newline='') as csvfile:
    fieldnames = ['id', 'username', 'email', 'full_name', 'apps']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()
    
    for user in extracted_users:
        writer.writerow(user)

# Calculate and print the execution time
end_time = time.time()
execution_time = end_time - start_time
print(f"Execution time: {execution_time} seconds.")

核心优化点及实现方案 #

• 复用HTTP会话：使用requests.Session()维持连接池，避免每次请求重复建立TCP连接，减少网络开销。
• 并行处理用户应用查询：用多线程并行发起1100个用户的appLinks请求，替代串行遍历，大幅压缩总耗时（建议并发数设置为10-20，平衡效率与限流风险）。
• 消除重复请求：仅获取一次用户-应用关联字典app_links，同时供生成用户数据和应用用户映射使用，避免两次遍历用户的冗余API调用。
• 优化数据存储格式：在app_links中直接存储应用名称列表，省去字符串拼接与拆分的额外计算开销。

优化后脚本 #

import requests
import csv
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

# 创建复用的HTTP会话
def create_session(okta_api_token):
    session = requests.Session()
    session.headers.update({
        'Authorization': f'SSWS {okta_api_token}',
        'Accept': 'application/json',
        'Content-Type': 'application/json'
    })
    return session

def get_all_users(domain, session):
    okta_url = f'https://{domain}.okta.com'
    url = f'{okta_url}/api/v1/users'
    
    all_users = []
    while url:
        response = session.get(url)
        response.raise_for_status()
        users = response.json()
        all_users += users
        
        links = response.links
        url = links.get('next', {}).get('url')
    
    return all_users

def fetch_user_apps(domain, session, user_id):
    okta_url = f'https://{domain}.okta.com'
    url = f'{okta_url}/api/v1/users/{user_id}/appLinks'
    
    apps = []
    while url:
        response = session.get(url)
        response.raise_for_status()
        app_links_data = response.json()
        apps += app_links_data
        
        links = response.links
        url = links.get('next', {}).get('url')
    
    return user_id, [app['label'] for app in apps]

def build_app_links(domain, okta_api_token, users):
    app_links = {}
    session = create_session(okta_api_token)
    
    # 控制并发数，避免触发Okta限流
    with ThreadPoolExecutor(max_workers=15) as executor:
        futures = [executor.submit(fetch_user_apps, domain, session, user['id']) 
                   for user in users if user.get('id')]
        
        for future in as_completed(futures):
            user_id, app_names = future.result()
            app_links[user_id] = app_names
    
    return app_links

def extract_user_data(users, app_links):
    user_data = []
    for user in users:
        user_id = user.get('id', 'N/A')
        profile = user.get('profile', {})
        user_name = profile.get('login', 'N/A')
        email = profile.get('email', 'N/A')
        first_name = profile.get('firstName', 'N/A')
        last_name = profile.get('lastName', 'N/A')
        full_name = f"{first_name} {last_name}"
        
        apps = ', '.join(app_links.get(user_id, []))
        user_data.append({
            'id': user_id,
            'username': user_name,
            'email': email,
            'full_name': full_name,
            'apps': apps
        })
    return user_data

def extract_app_user_data(users, app_links):
    app_user_data = {}
    for user in users:
        user_id = user.get('id')
        if not user_id:
            continue
        
        app_names = app_links.get(user_id, [])
        user_login = user.get('profile', {}).get('login', 'N/A')
        
        for app in app_names:
            if app not in app_user_data:
                app_user_data[app] = []
            app_user_data[app].append(user_login)
    
    return app_user_data

# 主执行逻辑
if __name__ == "__main__":
    domain = os.environ['OKTA_DOMAIN']
    okta_api_token = os.environ['OKTA_API_TOKEN']
    
    start_time = time.time()
    
    # 获取所有用户列表
    session = create_session(okta_api_token)
    users = get_all_users(domain, session)
    
    # 并行构建用户-应用关联字典
    app_links = build_app_links(domain, okta_api_token, users)
    
    # 生成两种输出数据
    extracted_users = extract_user_data(users, app_links)
    app_user_data = extract_app_user_data(users, app_links)
    
    # 保存用户数据到CSV
    with open('user_data.csv', 'w', newline='') as csvfile:
        fieldnames = ['id', 'username', 'email', 'full_name', 'apps']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(extracted_users)
    
    # 保存应用-用户映射到CSV
    with open('app_user_data.csv', 'w', newline='') as csvfile:
        fieldnames = ['app', 'users']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for app, user_logins in app_user_data.items():
            writer.writerow({'app': app, 'users': ', '.join(user_logins)})
    
    execution_time = time.time() - start_time
    print(f"Execution time: {execution_time:.2f} seconds.")

额外建议 #

• 若Okta API存在严格速率限制，可在fetch_user_apps中添加重试机制（如使用tenacity库），避免临时限流导致脚本中断。
• 根据实际网络环境和Okta平台的限流规则，调整ThreadPoolExecutor的max_workers参数，找到最优并发数。

内容的提问来源于stack exchange，提问作者rmrj