如何在CosmosDB SDK(.NET)中获取非迭代Linq查询的诊断信息

针对单结果Linq查询（Count()、Sum()、Single()等）的诊断捕获需求，结合你已有的自定义QueryProvider/Queryable，可通过以下方式实现统一拦截，无需逐个修改查询方式：

核心思路 #

绕开原生QueryProvider的直接执行逻辑，将单结果查询强制转换为FeedIterator执行，通过FeedResult获取诊断数据后，再提取对应单结果返回，以此保留RequestCharge等关键信息。

具体实现步骤 #

1. 重写自定义QueryProvider的执行方法 #

在自定义QueryProvider中，拦截同步/异步的Execute方法，替换为FeedIterator执行逻辑：

public class CustomCosmosQueryProvider : IQueryProvider
{
    private readonly IQueryProvider _originalProvider;
    private readonly Container _container;

    public CustomCosmosQueryProvider(IQueryProvider originalProvider, Container container)
    {
        _originalProvider = originalProvider;
        _container = container;
    }

    public TResult Execute<TResult>(Expression expression)
    {
        var query = _originalProvider.CreateQuery<TResult>(expression);
        var (result, requestCharge) = ExecuteWithDiagnostics<TResult>(query);
        
        // 统一记录诊断信息（示例：写入日志系统）
        Console.WriteLine($"Query RequestCharge: {requestCharge:F2} RU");
        
        return result;
    }

    public async Task<TResult> ExecuteAsync<TResult>(Expression expression, CancellationToken cancellationToken = default)
    {
        var query = _originalProvider.CreateQuery<TResult>(expression);
        var (result, requestCharge) = await ExecuteWithDiagnosticsAsync<TResult>(query, cancellationToken);
        
        Console.WriteLine($"Async Query RequestCharge: {requestCharge:F2} RU");
        return result;
    }
}

2. 封装通用诊断捕获逻辑 #

编写通用方法处理不同类型的单结果查询，从FeedResult中提取结果和RequestCharge：

private (TResult Result, double RequestCharge) ExecuteWithDiagnostics<TResult>(IQueryable<TResult> query)
{
    var cosmosQuery = query as CosmosQueryable<TResult> ?? 
                      throw new InvalidOperationException("Query is not a CosmosDB query");
    
    var queryDefinition = cosmosQuery.ToQueryDefinition();
    using var iterator = _container.GetItemQueryIterator<TResult>(queryDefinition);
    
    // 单结果查询仅返回一页，直接读取第一页
    var feedResult = iterator.ReadNextAsync().GetAwaiter().GetResult();
    
    // 根据返回类型提取结果
    TResult result = default;
    if (typeof(TResult) == typeof(long) || typeof(TResult) == typeof(int) || typeof(TResult) == typeof(decimal))
    {
        // 适配Count/Sum等数值型结果
        result = feedResult.FirstOrDefault();
    }
    else if (typeof(TResult).IsClass && typeof(TResult) != typeof(string))
    {
        // 适配Single/SingleOrDefault等实体结果
        result = feedResult.SingleOrDefault();
    }
    // 可扩展更多类型适配（如bool、string等）
    
    var totalRequestCharge = feedResult.Diagnostics.GetTotalRequestCharge();
    return (result, totalRequestCharge);
}

// 异步版本
private async Task<(TResult Result, double RequestCharge)> ExecuteWithDiagnosticsAsync<TResult>(IQueryable<TResult> query, CancellationToken cancellationToken)
{
    var cosmosQuery = query as CosmosQueryable<TResult> ?? 
                      throw new InvalidOperationException("Query is not a CosmosDB query");
    
    var queryDefinition = cosmosQuery.ToQueryDefinition();
    using var iterator = _container.GetItemQueryIterator<TResult>(queryDefinition);
    
    var feedResult = await iterator.ReadNextAsync(cancellationToken);
    
    TResult result = default;
    if (typeof(TResult) == typeof(long) || typeof(TResult) == typeof(int) || typeof(TResult) == typeof(decimal))
    {
        result = feedResult.FirstOrDefault();
    }
    else if (typeof(TResult).IsClass && typeof(TResult) != typeof(string))
    {
        result = feedResult.SingleOrDefault();
    }
    
    var totalRequestCharge = feedResult.Diagnostics.GetTotalRequestCharge();
    return (result, totalRequestCharge);
}

3. 关联自定义Queryable #

确保你的自定义Queryable使用上述CustomCosmosQueryProvider，所有通过该Queryable发起的Linq查询（包括单结果类型）都会自动被拦截，无需修改原有查询代码。

关键注意事项 #

• 类型适配：需根据实际使用的单结果查询类型（如CountAsync返回Task<long>，SumAsync<int>返回Task<int>）完善结果提取逻辑，避免类型转换错误。
• 异步覆盖：必须实现异步版本的ExecuteAsync方法，覆盖CountAsync、SingleAsync等常用异步查询场景。
• 诊断扩展：除了RequestCharge，还可从FeedResult.Diagnostics中提取查询耗时、重试次数、响应状态码等更多诊断数据，按需记录。

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