Python中并发调用Azure KeyVault set_secret API优化咨询

嘿，你的问题我太懂了——串行调用API处理几千个机密，光等网络往返就能耗掉几十分钟，换成异步并发绝对能把时间砍到几分钟以内。而且你完全不用手动搞多进程/线程，用Azure官方的异步SDK就能搞定，我给你一步步讲清楚：

第一步：替换成支持异步的新SDK #

你现在用的azure-keyvault是旧版SDK，已经被官方弃用了，而且没有原生异步支持。我们换成最新的异步版KeyVault机密SDK，它专门针对异步场景做了优化：

先安装依赖包：

pip install azure-keyvault-secrets[aio] azure-identity

第二步：异步实现核心逻辑 #

新SDK提供了SecretClient的异步版本（来自azure.keyvault.secrets.aio），配合异步凭据（比如ClientSecretCredential的异步实现），可以轻松并发发起API请求。核心思路是用asyncio.gather批量执行所有写入任务，而不是串行等待每个请求完成。

完整代码示例 #

import asyncio
import json
import sys
import os
from azure.keyvault.secrets.aio import SecretClient
from azure.identity.aio import ClientSecretCredential

# 从环境变量读取配置（推荐，避免硬编码敏感信息）
TENANT_ID = os.getenv("AZURE_TENANT_ID")
CLIENT_ID = os.getenv("AZURE_CLIENT_ID")
CLIENT_SECRET = os.getenv("AZURE_CLIENT_SECRET")
KV_URL = os.getenv("KEYVAULT_URL")
BASE_ENV = os.getenv("BASE_ENV")
BASE_FUNCTION = os.getenv("BASE_FUNCTION")

async def write_secrets():
    # 初始化异步凭据和KeyVault客户端
    credential = ClientSecretCredential(
        tenant_id=TENANT_ID,
        client_id=CLIENT_ID,
        client_secret=CLIENT_SECRET
    )
    client = SecretClient(vault_url=KV_URL, credential=credential)

    data = sys.stdin.read()
    if not data:
        print("OK nothing to do.")
        # 记得关闭异步资源
        await credential.close()
        await client.close()
        return

    print(f"Writing outputs to keyvault {KV_URL}")
    outputs = json.loads(data)

    # 准备所有异步写入任务
    tasks = []
    secret_names = []
    for key_name, value_data in outputs.items():
        secret_name = f"{BASE_ENV}-{BASE_FUNCTION}-{key_name}".replace("_", "-")
        secret_value = str(value_data["value"])
        print(f"Preparing to store: {secret_name} = {secret_value}")
        # 将每个写入操作包装成异步任务
        tasks.append(client.set_secret(secret_name, secret_value))
        secret_names.append(secret_name)

    # 并发执行所有任务，return_exceptions=True允许其他任务继续执行即使某个失败
    results = await asyncio.gather(*tasks, return_exceptions=True)

    # 统计成功/失败情况
    success_count = 0
    failed_secrets = []
    for secret_name, result in zip(secret_names, results):
        if isinstance(result, Exception):
            print(f"Error storing secret {secret_name}: {str(result)}")
            failed_secrets.append(secret_name)
        else:
            success_count += 1

    print(f"\n{success_count} secrets written successfully.")
    if failed_secrets:
        print(f"Failed to write {len(failed_secrets)} secrets: {', '.join(failed_secrets)}")
        # 如果需要失败就退出脚本，取消下面的注释
        # exit(os.EX_SOFTWARE)

    # 清理异步资源
    await credential.close()
    await client.close()

if __name__ == "__main__":
    asyncio.run(write_secrets())

关键细节说明 #

1. 为什么不用多客户端？
   异步版SecretClient内部已经维护了高效的连接池，单个客户端就能处理大量并发请求，完全不需要手动创建多个实例——反而多个实例会浪费资源。

2. 并发控制（避免KeyVault限流）
   Azure KeyVault有默认的速率限制（机密操作每30秒最多1000个请求），如果一次性发起4000个请求可能会触发限流。你可以改成分批处理，比如每50个一批，加个短暂延迟：

   # 修改任务执行部分，分批处理
   batch_size = 50
   success_count = 0
   failed_secrets = []
   
   for i in range(0, len(tasks), batch_size):
       batch_tasks = tasks[i:i+batch_size]
       batch_names = secret_names[i:i+batch_size]
       batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
       
       for name, res in zip(batch_names, batch_results):
           if isinstance(res, Exception):
               print(f"Error storing {name}: {str(res)}")
               failed_secrets.append(name)
           else:
               success_count += 1
       
       # 每批结束后短暂休眠，避免触发限流
       await asyncio.sleep(0.5)
   

3. 错误处理优化
   原来的代码遇到单个错误就直接退出，异步版本可以继续处理其他机密，最后统一统计失败情况，更适合批量场景。

效果对比 #

原来串行处理4000个机密需要50分钟，换成异步并发后，即使分批处理，总时间也能降到2-5分钟（取决于网络和KeyVault的响应速度），效率提升非常明显。

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