TEAMModelOS/TEAMModelOS.SDK/Module/AzureCosmosDBV3/AzureCosmosDBV3Repository.cs

using Microsoft.Azure.Cosmos;
using Microsoft.Azure.Cosmos.Linq;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;
using System.Net;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;
using TEAMModelOS.SDK.Context.Attributes.Azure;
using TEAMModelOS.SDK.Context.Exception;
using TEAMModelOS.SDK.Helper.Common.ReflectorExtensions;
using TEAMModelOS.SDK.Module.AzureCosmosDB.Configuration;

namespace TEAMModelOS.SDK.Module.AzureCosmosDBV3
{
    public class AzureCosmosDBV3Repository : IAzureCosmosDBV3Repository, IDisposable
    {
        private CosmosClient CosmosClient { get; set; }
        /// <summary>
        /// 线程安全的dict类型
        /// </summary>
        private Dictionary<string, Container> DocumentCollectionDict { get; set; } = new Dictionary<string, Container>();

        private string DatabaseId { get; set; }
        private int CollectionThroughput { get; set; }
        private Database database = null;

        private string[] ScanModel { get; set; }

        public AzureCosmosDBV3Repository(AzureCosmosDBOptions options)
        {
            try
            {
                if (!string.IsNullOrEmpty(options.ConnectionString))
                {
                    CosmosClient = CosmosDBV3ClientSingleton.getInstance(options.ConnectionString, options.ConnectionKey).GetCosmosDBClient();

                }
                else
                {
                    throw new BizException("请设置正确的AzureCosmosDB数据库配置信息！");
                }
                DatabaseId = options.Database;
                CollectionThroughput = options.CollectionThroughput;
                ScanModel = options.ScanModel;
                // InitializeDatabase().GetAwaiter().GetResult();
            }
            catch (CosmosException e)
            {
                Dispose(true);
                throw new BizException(e.Message, 500, e.StackTrace);
            }
        }


        public async Task InitializeDatabase()
        {
            try
            {
                database = await CosmosClient.CreateDatabaseIfNotExistsAsync(DatabaseId, CollectionThroughput);
                FeedIterator<ContainerProperties> resultSetIterator = database.GetContainerQueryIterator<ContainerProperties>();
                while (resultSetIterator.HasMoreResults)
                {
                    foreach (ContainerProperties container in await resultSetIterator.ReadNextAsync())
                    {
                        DocumentCollectionDict.TryAdd(container.Id, database.GetContainer(container.Id));
                    }
                }
                //获取数据库所有的表
                List<Type> types = ReflectorExtensions.GetAllTypeAsAttribute<CosmosDBAttribute>(ScanModel);
                foreach (Type type in types)
                {
                    string PartitionKey = GetPartitionKey(type);
                    string CollectionName = "";
                    int RU = 0;
                    IEnumerable<CosmosDBAttribute> attributes = type.GetCustomAttributes<CosmosDBAttribute>(true);
                    if (!string.IsNullOrEmpty(attributes.First<CosmosDBAttribute>().Name))
                    {
                        CollectionName = attributes.First<CosmosDBAttribute>().Name;
                    }
                    else
                    {
                        CollectionName = type.Name;
                    }
                    if (attributes.First<CosmosDBAttribute>().RU > 400)
                    {
                        RU = attributes.First<CosmosDBAttribute>().RU;
                    }
                    else
                    {
                        RU = CollectionThroughput;
                    }
                    //如果表存在于数据则检查RU是否变动，如果不存在则执行创建DocumentCollection
                    if (DocumentCollectionDict.TryGetValue(CollectionName, out Container collection))
                    { //更新RU

                        int? throughputResponse = await CosmosClient.GetDatabase(DatabaseId).GetContainer(collection.Id).ReadThroughputAsync();
                        if (throughputResponse < RU)
                        {
                            await CosmosClient.GetDatabase(DatabaseId).GetContainer(collection.Id).ReplaceThroughputAsync(RU);
                        }
                    }
                    else
                    {
                        ContainerProperties containerProperties = new ContainerProperties { Id = CollectionName };

                        if (!string.IsNullOrEmpty(PartitionKey))
                        {
                            containerProperties.PartitionKeyPath = "/" + PartitionKey;
                        }
                        if (RU > CollectionThroughput)
                        {
                            CollectionThroughput = RU;
                        }
                        Container containerWithConsistentIndexing = await database.CreateContainerIfNotExistsAsync(containerProperties, throughput: CollectionThroughput);
                        DocumentCollectionDict.TryAdd(CollectionName, containerWithConsistentIndexing);
                    }
                }
            }
            catch (CosmosException e)
            {
                throw new BizException(e.Message, 500, e.StackTrace);
            }
        }




        private string GetPartitionKey<T>()
        {
            Type type = typeof(T);
            return GetPartitionKey(type);
        }
        private string GetPartitionKey(Type type)
        {
            PropertyInfo[] properties = type.GetProperties();
            List<PropertyInfo> attrProperties = new List<PropertyInfo>();
            foreach (PropertyInfo property in properties)
            {
                if (property.Name.Equals("PartitionKey"))
                {
                    attrProperties.Add(property);
                    break;
                }
                object[] attributes = property.GetCustomAttributes(true);
                foreach (object attribute in attributes) //2.通过映射，找到成员属性上关联的特性类实例，
                {
                    if (attribute is PartitionKeyAttribute)
                    {
                        attrProperties.Add(property);
                    }
                }
            }
            if (attrProperties.Count <= 0)
            {
                throw new BizException(type.Name + "has no PartitionKey !");
            }
            else
            {
                if (attrProperties.Count == 1)
                {
                    return attrProperties[0].Name;
                }
                else { throw new BizException("PartitionKey can only be single!"); }
            }
        }


        private async Task<Container> InitializeCollection<T>()
        {
            Type type = typeof(T);
            string partitionKey = GetPartitionKey<T>();
            string CollectionName;
            IEnumerable<CosmosDBAttribute> attributes = type.GetCustomAttributes<CosmosDBAttribute>(true);
            if (!string.IsNullOrEmpty(attributes.First<CosmosDBAttribute>().Name))
            {
                CollectionName = attributes.First<CosmosDBAttribute>().Name;
            }
            else
            {
                CollectionName = type.Name;
            }
            return await InitializeCollection(CollectionName, partitionKey);
        }

        private async Task<Container> InitializeCollection(string CollectionName, string PartitionKey)
        {
            /////内存中已经存在这个表则直接返回
            if (DocumentCollectionDict.TryGetValue(CollectionName, out Container DocumentCollection))
            {
                return DocumentCollection;
            }///如果没有则尝试默认创建
            else
            {
                ContainerProperties containerProperties = new ContainerProperties { Id = CollectionName };
                if (!string.IsNullOrEmpty(PartitionKey))
                {
                    containerProperties.PartitionKeyPath = "/" + PartitionKey;
                }
                Container containerWithConsistentIndexing = await database.CreateContainerIfNotExistsAsync(containerProperties, throughput: CollectionThroughput);
                DocumentCollectionDict.TryAdd(CollectionName, containerWithConsistentIndexing);
                return containerWithConsistentIndexing;
            }
        }


        public async Task DeleteAll<T>(List<KeyValuePair<string, string>> ids) where T : ID
        {
            string partitionKey = GetPartitionKey<T>();
            await Task.Run(() => Parallel.ForEach(ids, (item) =>
            {
                Task.WaitAll(DeleteAsync<T>(item.Value, item.Key));
            }));
        }

        public async Task DeleteAll<T>(List<T> entities) where T : ID
        {
            string partitionKey = GetPartitionKey<T>();
            Type type = typeof(T);
            await Task.Run(() => Parallel.ForEach(entities, (item) =>
            {
                object o = type.GetProperty(partitionKey).GetValue(item, null);
                Task.WaitAll(DeleteAsync<T>(item.id, o.ToString()));
            }));
        }
        public async Task<T> DeleteAsync<T>(string id, string pk) where T : ID
        {
            Container container = await InitializeCollection<T>();
            ItemResponse<T> response = await container.DeleteItemAsync<T>(id: id, partitionKey: new PartitionKey(pk));
            return response.Resource;
        }

        public async Task<T> DeleteAsync<T>(T entity, string pk) where T : ID
        {
            Container container = await InitializeCollection<T>();
            string partitionKey = GetPartitionKey<T>();
            Type type = typeof(T);
            object o = type.GetProperty(partitionKey).GetValue(entity, null);
            ItemResponse<T> response = await container.DeleteItemAsync<T>(id: entity.id, partitionKey: new PartitionKey(o.ToString()));
            return response.Resource;

        }
        //public async Task<T> DeleteAsync<T>(string id) where T : ID
        //{
        //    Container container = await InitializeCollection<T>();
        //    ItemResponse<T> response = await container.DeleteItemAsync<T>(id: id, partitionKey: new PartitionKey(GetPartitionKey<T>()));
        //    return response.Resource;
        //}

        public async Task<List<T>> FindAll<T>() where T : ID
        {
            Container container = await InitializeCollection<T>();
            return await ResultsFromFeedIterator(container.GetItemQueryIterator<T>());

        }
        private async Task<List<T>> ResultsFromFeedIterator<T>(FeedIterator<T> query, int? maxItemCount = null)
        {
            List<T> results = new List<T>();
            while (query.HasMoreResults)
            {
                foreach (T t in await query.ReadNextAsync())
                {
                    results.Add(t);
                    if (results.Count == maxItemCount)
                    {
                        return results;
                    }
                }
            }
            return results;
        }
        private async Task<List<T>> ResultsFromFeedIterator<T>(FeedIterator<T> query, Func<List<T>, Task> batchAction, int itemsPerPage)
        {
            List<T> results = new List<T>();
            while (query.HasMoreResults)
            {
                if (results.Count() >= itemsPerPage)
                {
                    await batchAction(results);
                    results.Clear();
                }

                results.AddRange(await query.ReadNextAsync());
            }

            if (results.Count() > 0)
            {
                await batchAction(results);
                results.Clear();
            }

            return results;
        }
        public async Task<List<dynamic>> FindByDict(string CollectionName, Dictionary<string, object> dict, int itemsPerPage = -1, int? maxItemCount = null, string partitionKey = null)
        {
            if (DocumentCollectionDict.TryGetValue(CollectionName, out Container container))
            {
                //StringBuilder sql = new StringBuilder("select value(c) from c");
                //SQLHelper.GetSQL(dict, ref sql);
                //CosmosDbQuery cosmosDbQuery = new CosmosDbQuery
                //{
                //    QueryText = sql.ToString()

                //};
                StringBuilder sql = new StringBuilder("select value(c) from c");
                CosmosDbQuery cosmosDbQuery = SQLHelperParametric.GetSQL(dict, sql);
                QueryRequestOptions queryRequestOptions = GetDefaultQueryRequestOptions(itemsPerPage: GetEffectivePageSize(itemsPerPage, maxItemCount));
                FeedIterator<dynamic> query = container.GetItemQueryIterator<dynamic>(queryDefinition: cosmosDbQuery.CosmosQueryDefinition, requestOptions: queryRequestOptions);
                return await ResultsFromFeedIterator(query, maxItemCount);
            }
            else
            {
                throw new BizException("CollectionName named:" + CollectionName + " dose not exsit in Database!");
            }

        }

        public async Task<List<dynamic>> FindCountByDict(string CollectionName, Dictionary<string, object> dict, int itemsPerPage = -1, int? maxItemCount = null, string partitionKey = null)
        {
            if (DocumentCollectionDict.TryGetValue(CollectionName, out Container container))
            {
                //StringBuilder sql = new StringBuilder("select  value count(c)  from c");
                //SQLHelper.GetSQL(dict, ref sql);
                //CosmosDbQuery cosmosDbQuery = new CosmosDbQuery
                //{
                //    QueryText = sql.ToString()

                //};
                StringBuilder sql = new StringBuilder("select  value count(c)  from c");
                CosmosDbQuery cosmosDbQuery = SQLHelperParametric.GetSQL(dict, sql);
                QueryRequestOptions queryRequestOptions = GetDefaultQueryRequestOptions(itemsPerPage: GetEffectivePageSize(itemsPerPage, maxItemCount));
                FeedIterator<dynamic> query = container.GetItemQueryIterator<dynamic>(queryDefinition: cosmosDbQuery.CosmosQueryDefinition, requestOptions: queryRequestOptions);
                return await ResultsFromFeedIterator(query, maxItemCount);
            }
            else
            {
                throw new BizException("CollectionName named:" + CollectionName + " dose not exsit in Database!");
            }
        }

        public async Task<List<T>> FindByParams<T>(Dictionary<string, object> dict, int itemsPerPage = -1, int? maxItemCount = null, string partitionKey = null) where T : ID
        {
            return await FindByDict<T>(dict, itemsPerPage, maxItemCount, partitionKey);
        }
        public async Task<List<T>> FindByDict<T>(Dictionary<string, object> dict, int itemsPerPage = -1, int? maxItemCount = null, string partitionKey = null) where T : ID
        {

            StringBuilder sql = new StringBuilder("select value(c) from c");
            SQLHelper.GetSQL(dict, ref sql);
            CosmosDbQuery cosmosDbQuery = new CosmosDbQuery
            {
                QueryText = sql.ToString()

            };
            QueryRequestOptions queryRequestOptions = GetDefaultQueryRequestOptions(itemsPerPage: GetEffectivePageSize(itemsPerPage, maxItemCount));
            return await ResultsFromQueryAndOptions<T>(cosmosDbQuery, queryRequestOptions);
        }

        private async Task<List<T>> ResultsFromQueryAndOptions<T>(CosmosDbQuery cosmosDbQuery, QueryRequestOptions queryOptions, int? maxItemCount = null)
        {
            Container container = await InitializeCollection<T>();
            FeedIterator<T> query = container.GetItemQueryIterator<T>(
                queryDefinition: cosmosDbQuery.CosmosQueryDefinition,
                requestOptions: queryOptions);

            return await ResultsFromFeedIterator(query, maxItemCount);
        }
        private int GetEffectivePageSize(int itemsPerPage, int? maxItemCount)
        {
            return itemsPerPage == -1 ? maxItemCount ?? itemsPerPage : Math.Min(maxItemCount ?? itemsPerPage, itemsPerPage);
        }
        private QueryRequestOptions GetDefaultQueryRequestOptions(int? itemsPerPage = null,
        int? maxBufferedItemCount = null,
        int? maxConcurrency = null)
        {
            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
            {
                MaxItemCount = itemsPerPage == -1 ? 1000 : itemsPerPage,
                MaxBufferedItemCount = maxBufferedItemCount ?? 100,
                MaxConcurrency = maxConcurrency ?? 50
            };

            return queryRequestOptions;
        }
        private async Task<List<T>> ResultsFromQueryAndOptions<T>(CosmosDbQuery cosmosDbQuery, Func<List<T>, Task> batchAction, QueryRequestOptions queryOptions)
        {
            Container container = await InitializeCollection<T>();
            FeedIterator<T> query = container.GetItemQueryIterator<T>(
                queryDefinition: cosmosDbQuery.CosmosQueryDefinition,
                requestOptions: queryOptions);

            return await ResultsFromFeedIterator(query, batchAction, queryOptions.MaxItemCount ?? 0);
        }



        private QueryRequestOptions GetQueryRequestOptions(int itemsPerPage)
        {
            QueryRequestOptions queryRequestOptions = new QueryRequestOptions
            {
                MaxItemCount = itemsPerPage
            };

            return queryRequestOptions;
        }
        public async Task<List<T>> FindLinq<T>(Expression<Func<T, bool>> query = null, Expression<Func<T, object>> order = null, bool isDesc = false, int itemsPerPage = -1, int? maxItemCount = null) where T : ID
        {
            //QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(itemsPerPage);
            QueryRequestOptions queryRequestOptions = GetDefaultQueryRequestOptions(itemsPerPage: GetEffectivePageSize(itemsPerPage, maxItemCount));
            FeedIterator<T> feedIterator;
            Container container = await InitializeCollection<T>();

            if (query == null)
            {
                if (order != null)
                {
                    if (isDesc)
                    {
                        feedIterator = container
                       .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions).OrderByDescending(order)
                       .ToFeedIterator();
                    }
                    else
                    {
                        feedIterator = container
                       .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions).OrderBy(order)
                       .ToFeedIterator();
                    }
                }
                else
                {
                    feedIterator = container
                   .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions)
                   .ToFeedIterator();
                }
            }
            else
            {
                if (order != null)
                {
                    if (isDesc)
                    {
                        feedIterator = container
                        .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions)
                        .Where(query).OrderByDescending(order)
                        .ToFeedIterator();
                    }
                    else
                    {
                        feedIterator = container
                    .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions)
                    .Where(query).OrderBy(order)
                    .ToFeedIterator();
                    }
                }
                else
                {
                    feedIterator = container
                    .GetItemLinqQueryable<T>(requestOptions: queryRequestOptions)
                    .Where(query)
                    .ToFeedIterator();
                }
            }
            return await ResultsFromFeedIterator<T>(feedIterator);
        }

        public async Task<List<T>> FindSQL<T>(string sql, Dictionary<string, object> Parameters = null, int itemsPerPage = -1, int? maxItemCount = null) where T : ID
        {
            Container container = await InitializeCollection<T>();
            QueryRequestOptions queryOptions = GetQueryRequestOptions(GetEffectivePageSize(itemsPerPage, maxItemCount));
            if (Parameters != null)
            {
                CosmosDbQuery cosmosDbQuery = new CosmosDbQuery
                {
                    QueryText = sql,
                    Parameters = Parameters
                };
                FeedIterator<T> feedIterator = container
                .GetItemQueryIterator<T>(cosmosDbQuery.CosmosQueryDefinition, requestOptions: queryOptions);
                return await ResultsFromFeedIterator(feedIterator);
            }
            else
            {
                QueryDefinition queryDefinition = new QueryDefinition(sql);
                return await ResultsFromFeedIterator<T>(container.GetItemQueryIterator<T>(queryDefinition));
            }
        }

        //public async Task<List<T>> FindSQL<T>(string sql, bool isPK) where T : ID
        //{
        //    Container container = await InitializeCollection<T>();
        //    QueryDefinition queryDefinition = new QueryDefinition(sql);
        //    return await ResultsFromFeedIterator<T>(container.GetItemQueryIterator<T>(queryDefinition));
        //}

        public async Task<T> ReplaceObject<T>(T entity) where T : ID
        {
            Container container = await InitializeCollection<T>();
            ItemResponse<T> response = await container.ReplaceItemAsync(item: entity, id: entity.id);
            if (response.StatusCode.Equals(HttpStatusCode.OK))
            {
                return response.Resource;
            }
            else { throw new BizException("error"); }

        }

        //public async Task<T> ReplaceObject<T>(T entity, string key, string partitionKey) where T : ID
        //{
        //    Container container = await InitializeCollection<T>();
        //    ItemResponse<T> response = await container.ReplaceItemAsync(item: entity, id: entity.id);
        //    if (response.StatusCode.Equals(HttpStatusCode.OK))
        //    {
        //        return response.Resource;
        //    }
        //    else { throw new BizException("error"); }

        //}

        public async Task<T> Save<T>(T entity) where T : ID
        {
            try
            {
                Container container = await InitializeCollection<T>();
                ItemResponse<T> response = await container.CreateItemAsync<T>(entity);
                return response.Resource;
            }
            catch (Exception e)
            {
                throw new BizException(e.Message);
            }
        }

        public async Task<List<T>> SaveAll<T>(List<T> enyites) where T : ID
        {

            await Task.Run(() => Parallel.ForEach(enyites, (item) =>
            {
                Task.WaitAll(Save(item));
            }));
            return enyites;
        }

        public async Task<T> Update<T>(T entity) where T : ID
        {
            Container container = await InitializeCollection<T>();
            ItemResponse<T> response = await container.UpsertItemAsync(entity);
            return response.Resource;
        }

        public async Task<List<T>> UpdateAll<T>(List<T> entities) where T : ID
        {
            await Task.Run(() => Parallel.ForEach(entities, (item) =>
            {
                Task.WaitAll(Update(item));
            }));
            return entities;
        }

        public void Dispose()
        {
            Dispose(true);
        }
        protected virtual void Dispose(bool disposing)
        {
            if (disposing)
            {
                CosmosClient?.Dispose();
            }
        }

        public async Task<T> FindById<T>(string id) where T : ID
        {
            Container container = await InitializeCollection<T>();
            CosmosDbQuery cosmosDbQuery = new CosmosDbQuery
            {
                QueryText = @"SELECT *
                            FROM    c 
                            WHERE   c.id = @id",
                Parameters = new Dictionary<string, object>
                 {

                    { "@id",id}
                }
            };
            FeedIterator<T> feedIterator = container
               .GetItemQueryIterator<T>(cosmosDbQuery.CosmosQueryDefinition);
            return (await ResultsFromFeedIterator(feedIterator)).SingleOrDefault();
        }

        public async Task<T> FindByIdPk<T>(string id, string pk) where T : ID
        {
            Container container = await InitializeCollection<T>();
            ItemResponse<T> response = await container.ReadItemAsync<T>(id: id, partitionKey: new PartitionKey(pk));
            return response.Resource;
        }


        public async Task<List<T>> FindByDictTest<T>(Dictionary<string, object> dict, int itemsPerPage = 1, int? maxItemCount = 1, string partitionKey = null) where T : ID
        {
            //Container container = await InitializeCollection<T>();
            StringBuilder sql = new StringBuilder("select value(c) from c");
            CosmosDbQuery cosmosDbQuery = SQLHelperParametric.GetSQL(dict, sql);
            QueryRequestOptions queryRequestOptions = GetDefaultQueryRequestOptions(itemsPerPage: GetEffectivePageSize(itemsPerPage, maxItemCount));

            //FeedIterator<T> feedIterator = container
            //   .GetItemQueryIterator<T>(cosmosDbQuery.CosmosQueryDefinition);
            // return (await ResultsFromFeedIterator(feedIterator)).SingleOrDefault();


            return await ResultsFromQueryAndOptions<T>(cosmosDbQuery, queryRequestOptions, maxItemCount);
        }
    }
}