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 System.Linq.Expressions;
using TEAMModelOS.SDK;
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;
            }
        }

        /// <summary>
        /// 取得当前容器指定分区键的Count数，支持SQL Where条件，不支持排序
        /// </summary>
        /// <param name="container"></param>
        /// <param name="partitionkey"></param>
        /// <param name="queryWhere"></param>
        /// <returns></returns>
        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;
        }

        public static async Task<List<Response>> DeleteItemsStreamAsync(this CosmosContainer container, List<string> ids, string partitionkey )
        {
            List<Response> responses = new List<Response>();

            foreach (var id in ids)
            {
                try
                {
                    responses.Add(await container.DeleteItemStreamAsync(id, new PartitionKey(partitionkey)));
                }
                catch 
                {
                    continue;
                }
            }
            return responses;
        }

        public static async Task<List<ItemResponse<T>>> DeleteItemsAsync<T>(this CosmosContainer container, List<string> ids, string partitionkey)
        {
            List<ItemResponse<T>> responses = new List<ItemResponse<T>>();

            foreach (var id in ids)
            {
                try
                {
                    responses.Add(await container.DeleteItemAsync<T>(id, new PartitionKey(partitionkey)));
                }
                catch
                {
                    continue;
                }
            }
            return responses;
        }
    }
}