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;
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;
using TEAMModelOS.SDK.Models.Cosmos.OpenEntity;

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;
        }
        public static async Task<CosmosDBResult<T>> GetList<T>(this CosmosContainer container, QueryDefinition  queryDefinition, string partitionkey, string? continuationToken = null, int? pageSize = null)
        {
            List<T> list = new List<T>();
            double RU = 0;
            if (string.IsNullOrWhiteSpace(partitionkey))
            {
                return (new CosmosDBResult<T> { list = list, ru = RU, continuationToken = continuationToken });
            }

            await foreach (var item in container.GetItemQueryStreamIterator(queryDefinition: queryDefinition, continuationToken: continuationToken,
                requestOptions: new QueryRequestOptions { MaxItemCount = pageSize, PartitionKey = new PartitionKey(partitionkey) }))
            {
                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())
                    {
                        list.Add(obj.ToObject<T>());
                    }
                }
                if (queryDefinition.QueryText.Contains(" distinct ", StringComparison.OrdinalIgnoreCase) || queryDefinition.QueryText.Contains(" order by ", StringComparison.OrdinalIgnoreCase))
                {
                    continuationToken = null;
                }
                else
                {
                    continuationToken = item.GetContinuationToken();
                }
                RU  +=item.RU();
                if (pageSize.HasValue && pageSize.Value >= 0 && list.Count >= pageSize)
                {
                    break;
                }
            }
            //记录日志，RU开销大于400（开发测试），1000（正式）
            return (new CosmosDBResult<T> { list = list, ru = RU, continuationToken = continuationToken }); ;
        }

        public static async Task<CosmosDBResult<T>> GetList<T>(this CosmosContainer container, string sql, string partitionkey, string? continuationToken = null, int? pageSize = null)
        {
            List<T> list = new List<T>();
            double RU = 0;
            if (string.IsNullOrWhiteSpace(partitionkey))
            {
                return (new CosmosDBResult<T> { list = list, ru = RU, continuationToken = continuationToken });
            }

            await foreach (var item in container.GetItemQueryStreamIterator(queryText: sql, continuationToken: continuationToken, 
                requestOptions: new QueryRequestOptions { MaxItemCount = pageSize, PartitionKey = new PartitionKey(partitionkey) }) )
            {
                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()) {
                        list.Add(obj.ToObject<T>());
                    }
                }
                if (sql.Contains(" distinct ", StringComparison.OrdinalIgnoreCase) || sql.Contains(" order by ", StringComparison.OrdinalIgnoreCase))
                {
                    continuationToken = null;
                }
                else
                {
                    continuationToken = item.GetContinuationToken();
                }
                RU = item.RU();
                if (pageSize.HasValue && pageSize.Value >= 0 && list.Count >= pageSize)
                {
                    break;
                }
            }
            //记录日志，RU开销大于400（开发测试），1000（正式）
            return (new CosmosDBResult<T> { list = list, ru = RU, continuationToken = continuationToken }); ;
        }
        public class CosmosDBResult<T>
        {
            public List<T>? list { get; set; }
            public string? continuationToken { get; set; }
            public double ru { get; set; }
        }
    }
}