diff --git a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs b/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs
index c277e3d507..a83758f2d9 100644
--- a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs
+++ b/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs
@@ -5,7 +5,7 @@
#pragma warning disable ASPIRECOSMOS001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
var db = builder.AddAzureCosmosDB("cosmos")
- .WithDatabase("db", database => database.Containers.Add(new("entries", "/Id")))
+ .WithDatabase("db", database => database.Containers.Add(new("entries", "/id")))
.RunAsPreviewEmulator(e => e.WithDataExplorer());
#pragma warning restore ASPIRECOSMOS001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
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 35d4b232c2..84e2031686 100644
--- a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj
+++ b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj
@@ -21,6 +21,7 @@
+
diff --git a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBEmulatorHealthCheck.cs b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBEmulatorHealthCheck.cs
deleted file mode 100644
index 19665026f8..0000000000
--- a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBEmulatorHealthCheck.cs
+++ /dev/null
@@ -1,61 +0,0 @@
-// Licensed to the .NET Foundation under one or more agreements.
-// The .NET Foundation licenses this file to you under the MIT license.
-
-using Microsoft.Azure.Cosmos;
-using Microsoft.Extensions.Diagnostics.HealthChecks;
-
-namespace Aspire.Hosting.Azure.CosmosDB;
-
-///
-/// This health check also creates default databases and containers for the Azure CosmosDB Emulator.
-///
-internal sealed class AzureCosmosDBEmulatorHealthCheck : IHealthCheck
-{
- private readonly Func _clientFactory;
- private readonly Func _databasesFactory;
- private bool _resourcesCreated;
-
- public AzureCosmosDBEmulatorHealthCheck(Func clientFactory, Func databasesFactory)
- {
- ArgumentNullException.ThrowIfNull(clientFactory);
- ArgumentNullException.ThrowIfNull(databasesFactory);
-
- _clientFactory = clientFactory;
- _databasesFactory = databasesFactory;
- }
-
- ///
- public async Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
- {
- try
- {
- var cosmosClient = _clientFactory();
-
- await cosmosClient.ReadAccountAsync().WaitAsync(cancellationToken).ConfigureAwait(false);
-
- // Create the databases and containers if they do not exist. This is only performed once.
- if (!_resourcesCreated)
- {
- var databases = _databasesFactory();
-
- foreach (var database in databases)
- {
- var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.Name, cancellationToken: cancellationToken).ConfigureAwait(false)).Database;
-
- foreach (var container in database.Containers)
- {
- await db.CreateContainerIfNotExistsAsync(container.Name, container.PartitionKeyPath, cancellationToken: cancellationToken).ConfigureAwait(false);
- }
- }
-
- _resourcesCreated = true;
- }
-
- return HealthCheckResult.Healthy();
- }
- catch (Exception ex)
- {
- return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
- }
- }
-}
diff --git a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
index 899f42b73b..8630d3fd58 100644
--- a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
+++ b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs
@@ -14,7 +14,6 @@
using Azure.Provisioning.KeyVault;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.DependencyInjection;
-using Microsoft.Extensions.Diagnostics.HealthChecks;
namespace Aspire.Hosting;
@@ -97,18 +96,32 @@ private static IResourceBuilder RunAsEmulator(this IResou
cosmosClient = CreateCosmosClient(connectionString);
});
+ builder.ApplicationBuilder.Eventing.Subscribe(builder.Resource, async (@event, ct) =>
+ {
+ if (cosmosClient is null)
+ {
+ throw new InvalidOperationException("CosmosClient is not initialized.");
+ }
+
+ await cosmosClient.ReadAccountAsync().WaitAsync(ct).ConfigureAwait(false);
+
+ foreach (var database in builder.Resource.Databases)
+ {
+ var db = (await cosmosClient.CreateDatabaseIfNotExistsAsync(database.Name, cancellationToken: ct).ConfigureAwait(false)).Database;
+
+ foreach (var container in database.Containers)
+ {
+ await db.CreateContainerIfNotExistsAsync(container.Name, container.PartitionKeyPath, cancellationToken: ct).ConfigureAwait(false);
+ }
+ }
+ });
+
// Use custom health check that also seeds the databases and containers
var healthCheckKey = $"{builder.Resource.Name}_check";
- builder.ApplicationBuilder.Services.AddHealthChecks().Add(
- new HealthCheckRegistration(
- name: healthCheckKey,
- new AzureCosmosDBEmulatorHealthCheck(
- () => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),
- builder.Resource.Databases.ToArray
- ),
- failureStatus: null,
- tags: null)
- );
+ builder.ApplicationBuilder.Services.AddHealthChecks().AddAzureCosmosDB(
+ sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),
+ name: healthCheckKey
+ );
builder.WithHealthCheck(healthCheckKey);