如何用Python/Spark动态调整Cosmos DB容器RUs？代码无效排查

原代码存在的问题 #

• 使用client.create_database(DATABASE_NAME)会尝试创建新数据库，若数据库已存在会抛出异常，导致后续容器操作无法正确指向目标数据库，应改用get_database_client获取已存在的数据库实例。
• 代码中读取和替换吞吐量的逻辑是伪代码，未正确执行，且替换后没有重新读取吞吐量进行验证。
• 未考虑权限因素：操作吞吐量需要账号拥有Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/throughput/write权限，权限不足会导致修改不生效。

正确实现代码 #

from azure.cosmos import CosmosClient

# 初始化客户端
URL = "你的Cosmos DB账户URL"
KEY = "你的账户密钥"
client = CosmosClient(URL, credential=KEY)

# 获取目标数据库和容器客户端（避免重复创建）
DATABASE_NAME = 'db_name'
database = client.get_database_client(DATABASE_NAME)
container_name = 'container_name'
container = database.get_container_client(container_name)

# 1. 读取初始吞吐量并记录
initial_offer = container.read_offer()
initial_ru = initial_offer.offer_throughput
print(f"初始吞吐量: {initial_ru} RUs")

try:
    # 2. 提高吞吐量到目标值（比如10000）
    updated_offer = container.replace_throughput(10000)
    print(f"已将吞吐量调整为: {updated_offer.offer_throughput} RUs")

    # 3. 执行PySpark数据插入任务
    # 这里插入你的PySpark数据处理代码
    # spark.read.load(...).write.format("cosmos.oltp").options(...).save()

finally:
    # 4. 任务完成后恢复初始吞吐量
    restored_offer = container.replace_throughput(initial_ru)
    print(f"已恢复吞吐量为初始值: {restored_offer.offer_throughput} RUs")

额外注意事项 #

• 吞吐量调整延迟：Cosmos DB调整吞吐量可能需要几秒到几十秒，若需要立即验证，可添加短暂等待后重新读取吞吐量。
• 异常处理：在数据插入和吞吐量调整过程中添加异常捕获，避免因任务失败导致吞吐量未恢复。
• 自动缩放吞吐量：若频繁调整，可考虑开启Cosmos DB的自动缩放功能，系统会根据负载自动调整RUs，但需注意成本差异。

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