如何在Entity Framework中获取每次操作后的Azure Cosmos DB请求单元(RU)

有几种简便的编程方式可以获取EF Cosmos DB操作的RU消耗，完全满足集成测试、基准测试这类场景的需求，本质都是从Cosmos返回的x-ms-request-charge响应头提取数值，下面是具体实现：

1. 使用EF Core拦截器（推荐，和EF集成度最高） #

通过自定义EF Core拦截器，在每次数据库操作后自动提取RU值，不需要修改业务代码，适合测试或全局监控场景。

首先创建拦截器类：

public class CosmosRuTrackingInterceptor : DbCommandInterceptor
{
    // 用于存储累计RU，测试场景可以直接通过这个属性获取结果
    public decimal TotalRequestCharge { get; private set; }

    private void ExtractRuFromCommand(DbCommand command)
    {
        if (command is CosmosDbCommand cosmosCommand)
        {
            if (cosmosCommand.ResponseHeaders.TryGetValue("x-ms-request-charge", out var ruString)
                && decimal.TryParse(ruString, out var ruValue))
            {
                TotalRequestCharge += ruValue;
                // 这里也可以直接写入日志、监控系统等
            }
        }
    }

    // 覆盖需要拦截的操作方法，根据业务场景选择对应方法
    public override InterceptionResult<DbDataReader> ReaderExecuting(DbCommand command, CommandEventData eventData, InterceptionResult<DbDataReader> result)
    {
        ExtractRuFromCommand(command);
        return base.ReaderExecuting(command, eventData, result);
    }

    public override InterceptionResult<int> NonQueryExecuting(DbCommand command, CommandEventData eventData, InterceptionResult<int> result)
    {
        ExtractRuFromCommand(command);
        return base.NonQueryExecuting(command, eventData, result);
    }

    public override InterceptionResult<object> ScalarExecuting(DbCommand command, CommandEventData eventData, InterceptionResult<object> result)
    {
        ExtractRuFromCommand(command);
        return base.ScalarExecuting(command, eventData, result);
    }
}

然后在DbContext中注册拦截器：

protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
{
    optionsBuilder
        .UseCosmos("你的连接字符串", "数据库名称")
        .AddInterceptors(new CosmosRuTrackingInterceptor());
}

在集成测试里，你可以实例化拦截器并获取累计RU：

[Fact]
public async Task VerifyOperationRuUsage()
{
    var ruInterceptor = new CosmosRuTrackingInterceptor();
    var dbContextOptions = new DbContextOptionsBuilder<YourDbContext>()
        .UseCosmos("测试环境连接字符串", "测试数据库")
        .AddInterceptors(ruInterceptor)
        .Options;

    using var context = new YourDbContext(dbContextOptions);
    // 执行测试操作
    await context.YourEntities.AddAsync(new YourEntity { Id = "test-id" });
    await context.SaveChangesAsync();

    // 断言RU消耗符合预期范围
    Assert.InRange(ruInterceptor.TotalRequestCharge, 0, 15);
}

2. 直接使用Cosmos SDK获取RU #

如果你的代码中直接使用CosmosClient而非仅通过EF操作，那么可以直接从请求响应对象中获取RequestCharge属性：

var cosmosClient = new CosmosClient("你的连接字符串");
var container = cosmosClient.GetContainer("数据库名称", "容器名称");

// 单条写入操作的RU
var createResponse = await container.CreateItemAsync(new YourEntity { Id = "1" });
decimal createRu = createResponse.RequestCharge;

// 查询操作的RU
var queryResponse = await container.GetItemQueryIterator<YourEntity>("SELECT * FROM c").ReadNextAsync();
decimal queryRu = queryResponse.RequestCharge;

3. 从HttpClient响应头手动提取（不推荐，需自定义HttpClient） #

如果需要绕开EF和Cosmos SDK的封装，直接通过HttpClient调用Cosmos REST API，可以从响应头中读取x-ms-request-charge：

using var httpClient = new HttpClient();
httpClient.DefaultRequestHeaders.Add("Authorization", "你的授权令牌");

var response = await httpClient.PostAsync("Cosmos REST API地址", new StringContent("请求体"));
if (response.Headers.TryGetValues("x-ms-request-charge", out var ruValues))
{
    decimal ruCharge = decimal.Parse(ruValues.First());
}

这种方法需要自己处理认证和请求序列化，一般不如前两种简便，仅在特殊场景下使用。

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