Skip to content

Commit

Permalink
Merge pull request #80 from AzureCosmosDB/feature/cosmos-error-reporting
Browse files Browse the repository at this point in the history 
Improving record level error reporting for Cosmos
  • Loading branch information
@bowencode
bowencode committed Jul 27, 2023
2 parents 5ba8a1f + 1cffd38 commit f458054
Showing 1 changed file with 62 additions and 29 deletions.
91 changes: 62 additions & 29 deletions Extensions/Cosmos/Cosmos.DataTransfer.CosmosExtension/CosmosDataSinkExtension.cs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
using System.ComponentModel.Composition;
using System.Diagnostics;
using System.Dynamic;
using System.Net;
using System.Text;
using System.Text.RegularExpressions;
using Cosmos.DataTransfer.Interfaces;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.Configuration;
Expand Down Expand Up @@ -81,6 +81,7 @@ public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfigurati
await CosmosExtensionServices.VerifyContainerAccess(container, settings.Container, logger, cancellationToken);

int addedCount = 0;
int inputCount = 0;

var timer = Stopwatch.StartNew();
void ReportCount(int i)
Expand All @@ -100,15 +101,17 @@ void ReportCount(int i)
var addTasks = batch.Select(item => AddItemAsync(container, item, settings.PartitionKeyPath ?? settings.PartitionKeyPaths?.FirstOrDefault(), settings.WriteMode, retry, logger, cancellationToken)).ToList();

var results = await Task.WhenAll(addTasks);
ReportCount(results.Sum());
ReportCount(results.Sum(i => i.ItemCount));
inputCount += results.Length;
}

logger.LogInformation("Added {AddedCount} total records in {TotalSeconds}s", addedCount, $"{timer.ElapsedMilliseconds / 1000.0:F2}");
}
if (addedCount != inputCount)
{
logger.LogWarning("Added {AddedCount} of {TotalCount} total records in {TotalSeconds}s", addedCount, inputCount, $"{timer.ElapsedMilliseconds / 1000.0:F2}");
throw new Exception($"Only {addedCount} of {inputCount} records were added to Cosmos");
}

private static string StripSpecialChars(string displayName)
{
return Regex.Replace(displayName, "[^\\w]", "", RegexOptions.Compiled);
logger.LogInformation("Added {AddedCount} total records in {TotalSeconds}s", addedCount, $"{timer.ElapsedMilliseconds / 1000.0:F2}");
}

private static AsyncRetryPolicy GetRetryPolicy(int maxRetryCount, int initialRetryDuration)
Expand All @@ -124,44 +127,68 @@ private static AsyncRetryPolicy GetRetryPolicy(int maxRetryCount, int initialRet
return retryPolicy;
}

private static Task<int> AddItemAsync(Container container, ExpandoObject item, string? partitionKeyPath, DataWriteMode mode, AsyncRetryPolicy retryPolicy, ILogger logger, CancellationToken cancellationToken)
private static Task<ItemResult> AddItemAsync(Container container, ExpandoObject item, string? partitionKeyPath, DataWriteMode mode, AsyncRetryPolicy retryPolicy, ILogger logger, CancellationToken cancellationToken)
{
logger.LogTrace("Adding item {Id}", GetPropertyValue(item, "id"));
var task = retryPolicy.ExecuteAsync(() =>
{
switch (mode)
{
case DataWriteMode.InsertStream:
ArgumentNullException.ThrowIfNull(partitionKeyPath, nameof(partitionKeyPath));
return container.CreateItemStreamAsync(CreateItemStream(item), new PartitionKey(GetPropertyValue(item, partitionKeyPath.TrimStart('/'))), cancellationToken: cancellationToken);
case DataWriteMode.Insert:
return container.CreateItemAsync(item, cancellationToken: cancellationToken);
case DataWriteMode.UpsertStream:
ArgumentNullException.ThrowIfNull(partitionKeyPath, nameof(partitionKeyPath));
return container.UpsertItemStreamAsync(CreateItemStream(item), new PartitionKey(GetPropertyValue(item, partitionKeyPath.TrimStart('/'))), cancellationToken: cancellationToken);
case DataWriteMode.Upsert:
return container.UpsertItemAsync(item, cancellationToken: cancellationToken);
}
string? id = GetPropertyValue(item, "id");
logger.LogTrace("Adding item {Id}", id);

throw new ArgumentOutOfRangeException(nameof(mode), $"Invalid data write mode specified: {mode}");
})
var task = retryPolicy.ExecuteAsync(() => PopulateItem(container, item, partitionKeyPath, mode, id, cancellationToken))
.ContinueWith(t =>
{
if (t.IsCompletedSuccessfully)
bool requestSucceeded = t.Result.IsSuccess;
if (t.IsCompletedSuccessfully && requestSucceeded)
{
return 1;
return t.Result;
}
if (t.IsFaulted)
{
logger.LogWarning(t.Exception, "Error adding record: {ErrorMessage}", t.Exception?.Message);
}
else if (!requestSucceeded)
{
logger.LogWarning(t.Exception, "Error adding record {Id}: {ErrorMessage}", t.Result.Id, t.Result.StatusCode);
return t.Result;
}
return 0;
return new ItemResult(null, HttpStatusCode.InternalServerError);
}, cancellationToken);
return task;
}

private static async Task<ItemResult> PopulateItem(Container container, ExpandoObject item, string? partitionKeyPath, DataWriteMode mode, string? itemId, CancellationToken cancellationToken)
{
HttpStatusCode? statusCode = null;
switch (mode)
{
case DataWriteMode.InsertStream:
ArgumentNullException.ThrowIfNull(partitionKeyPath, nameof(partitionKeyPath));
var insertMessage = await container.CreateItemStreamAsync(CreateItemStream(item), new PartitionKey(GetPropertyValue(item, partitionKeyPath.TrimStart('/'))), cancellationToken: cancellationToken);
statusCode = insertMessage.StatusCode;
break;
case DataWriteMode.Insert:
var insertResponse = await container.CreateItemAsync(item, cancellationToken: cancellationToken);
statusCode = insertResponse.StatusCode;
break;
case DataWriteMode.UpsertStream:
ArgumentNullException.ThrowIfNull(partitionKeyPath, nameof(partitionKeyPath));
var upsertMessage = await container.UpsertItemStreamAsync(CreateItemStream(item), new PartitionKey(GetPropertyValue(item, partitionKeyPath.TrimStart('/'))), cancellationToken: cancellationToken);
statusCode = upsertMessage.StatusCode;
break;
case DataWriteMode.Upsert:
var upsertResponse = await container.UpsertItemAsync(item, cancellationToken: cancellationToken);
statusCode = upsertResponse.StatusCode;
break;
}

if (statusCode == null)
{
throw new ArgumentOutOfRangeException(nameof(mode), $"Invalid data write mode specified: {mode}");
}

return new ItemResult(itemId, statusCode.Value);
}

private static MemoryStream CreateItemStream(ExpandoObject item)
{
var json = JsonConvert.SerializeObject(item);
Expand All @@ -177,5 +204,11 @@ public IEnumerable<IDataExtensionSettings> GetSettings()
{
yield return new CosmosSinkSettings();
}

public record ItemResult(string? Id, HttpStatusCode StatusCode)
{
public bool IsSuccess => StatusCode is HttpStatusCode.OK or HttpStatusCode.Created;
public int ItemCount => IsSuccess ? 1 : 0;
}
}
}

0 comments on commit f458054

Please sign in to comment.