diff --git a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj
index 0806b9a815e..a2aade08516 100644
--- a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj
+++ b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj
@@ -24,6 +24,7 @@
+
diff --git a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
index 782961c8b38..d754874a09d 100644
--- a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
+++ b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
@@ -14,6 +14,8 @@
using Azure.Provisioning.KeyVault;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Polly;
namespace Aspire.Hosting;
@@ -101,6 +103,8 @@ private static IResourceBuilder RunAsEmulator(this IResou
cosmosClient = CreateCosmosClient(connectionString);
});
+ var creationState = HealthCheckResult.Unhealthy("Waiting for databases and containers to be created");
+
builder.ApplicationBuilder.Eventing.Subscribe(builder.Resource, async (@event, ct) =>
{
if (cosmosClient is null)
@@ -108,28 +112,45 @@ private static IResourceBuilder RunAsEmulator(this IResou
throw new InvalidOperationException("CosmosClient is not initialized.");
}
- await cosmosClient.ReadAccountAsync().WaitAsync(ct).ConfigureAwait(false);
-
- foreach (var database in builder.Resource.Databases)
+ try
{
- var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.DatabaseName, cancellationToken: ct).ConfigureAwait(false)).Database;
+ await cosmosClient.ReadAccountAsync().WaitAsync(ct).ConfigureAwait(false);
- foreach (var container in database.Containers)
+ foreach (var database in builder.Resource.Databases)
{
- var containerProperties = container.ContainerProperties;
+ var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.DatabaseName, cancellationToken: ct).ConfigureAwait(false)).Database;
+
+ foreach (var container in database.Containers)
+ {
+ var containerProperties = container.ContainerProperties ?? new ContainerProperties
+ {
+ Id = container.ContainerName,
+ PartitionKeyPaths = container.PartitionKeyPaths
+ };
- await db.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: ct).ConfigureAwait(false);
+ await db.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: ct).ConfigureAwait(false);
+ }
}
+ creationState = HealthCheckResult.Healthy();
+ }
+ catch (Exception ex)
+ {
+ creationState = HealthCheckResult.Degraded("Could not create databases and containers", ex);
+ throw;
}
});
var healthCheckKey = $"{builder.Resource.Name}_check";
- builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureCosmosDB(
- sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),
- name: healthCheckKey
- );
+ var creationHealthCheckKey = $"{builder.Resource.Name}_databases";
+ builder.ApplicationBuilder.Services.AddHealthChecks()
+ .AddAzureCosmosDB(
+ sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),
+ name: healthCheckKey
+ )
+ .AddCheck(creationHealthCheckKey, () => creationState);
- builder.WithHealthCheck(healthCheckKey);
+ builder.WithHealthCheck(healthCheckKey)
+ .WithHealthCheck(creationHealthCheckKey);
if (configureContainer != null)
{
@@ -155,11 +176,37 @@ static CosmosClient CreateCosmosClient(string connectionString)
{
clientOptions.ConnectionMode = ConnectionMode.Gateway;
clientOptions.LimitToEndpoint = true;
+ clientOptions.CustomHandlers.Add(new CosmosClientRetryHandler());
}
return new CosmosClient(connectionString, clientOptions);
}
}
+
+ }
+ class CosmosClientRetryHandler : RequestHandler
+ {
+ private static ResiliencePipeline Pipeline { get; }
+ = new ResiliencePipelineBuilder()
+ .AddRetry(new()
+ {
+ MaxRetryAttempts = 10,
+ Delay = TimeSpan.FromMilliseconds(500),
+ BackoffType = DelayBackoffType.Constant,
+ ShouldHandle = new PredicateBuilder()
+ .Handle()
+ .HandleResult(static result => !result.IsSuccessStatusCode),
+ })
+ .Build();
+
+ public override async Task SendAsync(
+ RequestMessage request,
+ CancellationToken cancellationToken)
+ {
+ return await Pipeline
+ .ExecuteAsync(async ct => await base.SendAsync(request, ct).ConfigureAwait(false), cancellationToken)
+ .ConfigureAwait(false);
+ }
}
///