TEAMModelOS/TEAMModelOS.SDK/DI/AzureCosmos/AzureCosmosExtensions.cs

using Microsoft.Azure.Cosmos.Table;
using Microsoft.Azure.Cosmos.Table.Queryable;
using System;
using System.Collections;
using System.Collections.Generic;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Linq;
using Azure;
using TEAMModelOS.SDK.DI.AzureCosmos.Inner;
using System.IO;
using TEAMModelOS.SDK.DI;
using System.Diagnostics;
using Azure.Cosmos;
using System.Text.Json;
using System.Net;
using TEAMModelOS.SDK.Context.Exception;
using System.Linq.Expressions;
using TEAMModelOS.SDK.Helper.Common.CollectionHelper;
using TEAMModelOS.SDK.Helper.Common.JsonHelper;
using TEAMModelOS.SDK.Extension;
using TEAMModelOS.SDK.Models;

namespace TEAMModelOS.SDK.DI
{
    public static class AzureCosmosExtensions
    {    
        public static double RU(this Response response)
        {
            try
            {
                response.Headers.TryGetValue("x-ms-request-charge", out var value);
                var ru = Convert.ToDouble(value);
                return ru;
            }
            catch
            {
                return 0;
            }
        }

        public static string GetContinuationToken(this Response response)
        {
            try
            {
                response.Headers.TryGetValue("x-ms-continuation", out var value);
                return value;
            }
            catch
            {
                return null;
            }
        }

        public static async Task<int> GetCount(this CosmosContainer container, string partitionkey, string queryWhere = "WHERE 1=1")
        {
            int totalCount = 0;
            await foreach (var item in container.GetItemQueryStreamIterator(
                queryText: $"SELECT VALUE COUNT(1) From c {queryWhere}",
                requestOptions: new QueryRequestOptions() { PartitionKey = new PartitionKey(partitionkey), MaxItemCount = -1 }))
            {
                using var json = await JsonDocument.ParseAsync(item.ContentStream);
                if (json.RootElement.TryGetProperty("_count", out JsonElement count) && count.GetUInt16() > 0)
                {
                    foreach (var obj in json.RootElement.GetProperty("Documents").EnumerateArray())
                    {
                        totalCount = obj.GetInt32();
                    }
                }
            }
            return totalCount;
        }
    }
}