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;
}
}
///
/// 取得当前容器指定分区键的Count数,支持SQL Where条件,不支持排序
///
///
///
///
///
public static async Task 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> DeleteItemsStreamAsync(this CosmosContainer container, List ids, string partitionkey )
{
List responses = new List();
foreach (var id in ids)
{
try
{
responses.Add(await container.DeleteItemStreamAsync(id, new PartitionKey(partitionkey)));
}
catch
{
continue;
}
}
return responses;
}
public static async Task>> DeleteItemsAsync(this CosmosContainer container, List ids, string partitionkey)
{
List> responses = new List>();
foreach (var id in ids)
{
try
{
responses.Add(await container.DeleteItemAsync(id, new PartitionKey(partitionkey)));
}
catch
{
continue;
}
}
return responses;
}
public static async Task> GetList(this CosmosContainer container, QueryDefinition queryDefinition, string partitionkey, string? continuationToken = null, int? pageSize = null)
{
List list = new List();
double RU = 0;
if (string.IsNullOrWhiteSpace(partitionkey))
{
return (new CosmosDBResult { 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());
}
}
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 { list = list, ru = RU, continuationToken = continuationToken }); ;
}
public static async Task> GetList(this CosmosContainer container, string sql, string partitionkey, string? continuationToken = null, int? pageSize = null)
{
List list = new List();
double RU = 0;
if (string.IsNullOrWhiteSpace(partitionkey))
{
return (new CosmosDBResult { 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());
}
}
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 { list = list, ru = RU, continuationToken = continuationToken }); ;
}
public class CosmosDBResult
{
public List? list { get; set; }
public string? continuationToken { get; set; }
public double ru { get; set; }
}
}
}