WaitFor for Cosmos DB (#5729)

* First pass of getting Cosmos going with Health Checks.

* Debugging persistent containers + cosmos + waitfor.

* Add Cosmos emulator functional test.

* Remove playground lifetime change.

* Create CosmosClient once.

Directory.Packages.props

@@ -70,6 +70,7 @@
     <PackageVersion Include="AspNetCore.HealthChecks.Azure.Storage.Blobs" Version="8.0.1" />
     <PackageVersion Include="AspNetCore.HealthChecks.Azure.Storage.Queues" Version="8.0.1" />
     <PackageVersion Include="AspNetCore.HealthChecks.AzureServiceBus" Version="8.0.1" />
+    <PackageVersion Include="AspNetCore.HealthChecks.CosmosDb" Version="8.0.1" />
     <PackageVersion Include="AspNetCore.HealthChecks.Kafka" Version="8.0.1" />
     <PackageVersion Include="AspNetCore.HealthChecks.MongoDb" Version="8.1.0" />
     <PackageVersion Include="AspNetCore.HealthChecks.MySql" Version="8.0.1" />

playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs

@@ -9,7 +9,7 @@
 
 builder.AddProject<Projects.CosmosEndToEnd_ApiService>("api")
        .WithExternalHttpEndpoints()
-       .WithReference(db);
+       .WithReference(db).WaitFor(db);
 
 #if !SKIP_DASHBOARD_REFERENCE
 // This project is only added in playground projects to support development/debugging

src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj

@@ -14,10 +14,13 @@
 
   <ItemGroup>
     <Compile Include="..\Shared\Cosmos\CosmosConstants.cs" Link="Shared\CosmosConstants.cs" />
+    <Compile Include="..\Shared\Cosmos\CosmosUtils.cs" Link="Shared\CosmosUtils.cs" />
   </ItemGroup>
 
   <ItemGroup>
     <ProjectReference Include="..\Aspire.Hosting.Azure\Aspire.Hosting.Azure.csproj" />
+    <PackageReference Include="Azure.Identity" />
+    <PackageReference Include="AspNetCore.HealthChecks.CosmosDb" />
     <PackageReference Include="Azure.Provisioning" />
     <PackageReference Include="Azure.Provisioning.CosmosDB" />
   </ItemGroup>

src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs

@@ -3,10 +3,14 @@
 
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Azure;
+using Aspire.Hosting.Azure.Cosmos;
+using Azure.Identity;
 using Azure.Provisioning;
 using Azure.Provisioning.CosmosDB;
 using Azure.Provisioning.KeyVaults;
 using Azure.ResourceManager.CosmosDB.Models;
+using Microsoft.Azure.Cosmos;
+using Microsoft.Extensions.DependencyInjection;
 using System.Diagnostics.CodeAnalysis;
 
 namespace Aspire.Hosting;
@@ -69,9 +73,52 @@ public static IResourceBuilder<AzureCosmosDBResource> AddAzureCosmosDB(this IDis
         };
 
         var resource = new AzureCosmosDBResource(name, configureConstruct);
+
+        CosmosClient? cosmosClient = null;
+
+        builder.Eventing.Subscribe<ConnectionStringAvailableEvent>(resource, async (@event, ct) =>
+        {
+            var connectionString = await resource.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
+
+            if (connectionString == null)
+            {
+                throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{resource.Name}' resource but the connection string was null.");
+            }
+
+            cosmosClient = CreateCosmosClient(connectionString);
+        });
+
+        var healthCheckKey = $"{name}_check";
+        builder.Services.AddHealthChecks().AddAzureCosmosDB(sp =>
+        {
+            return cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized.");
+        }, name: healthCheckKey);
+
         return builder.AddResource(resource)
                       .WithParameter(AzureBicepResource.KnownParameters.KeyVaultName)
-                      .WithManifestPublishingCallback(resource.WriteToManifest);
+                      .WithManifestPublishingCallback(resource.WriteToManifest)
+                      .WithHealthCheck(healthCheckKey);
+
+        static CosmosClient CreateCosmosClient(string connectionString)
+        {
+            var clientOptions = new CosmosClientOptions();
+            clientOptions.CosmosClientTelemetryOptions.DisableDistributedTracing = true;
+
+            if (Uri.TryCreate(connectionString, UriKind.Absolute, out var uri))
+            {
+                return new CosmosClient(uri.OriginalString, new DefaultAzureCredential(), clientOptions);
+            }
+            else
+            {
+                if (CosmosUtils.IsEmulatorConnectionString(connectionString))
+                {
+                    clientOptions.ConnectionMode = ConnectionMode.Gateway;
+                    clientOptions.LimitToEndpoint = true;
+                }
+
+                return new CosmosClient(connectionString, clientOptions);
+            }
+        }
     }
 
     /// <summary>

tests/Aspire.Hosting.Azure.Tests/AzureCosmosDBEmulatorFunctionalTests.cs

@@ -0,0 +1,59 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Aspire.Components.Common.Tests;
+using Aspire.Hosting.ApplicationModel;
+using Aspire.Hosting.Utils;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Aspire.Hosting.Azure.Tests;
+
+public class AzureCosmosDBEmulatorFunctionalTests(ITestOutputHelper testOutputHelper)
+{
+    [Fact]
+    [RequiresDocker]
+    public async Task VerifyWaitForOnCosmosDBEmulatorBlocksDependentResources()
+    {
+        // Cosmos can be pretty slow to spin up, lets give it plenty of time.
+        var cts = new CancellationTokenSource(TimeSpan.FromMinutes(5));
+        using var builder = TestDistributedApplicationBuilder.Create(testOutputHelper);
+
+        var healthCheckTcs = new TaskCompletionSource<HealthCheckResult>();
+        builder.Services.AddHealthChecks().AddAsyncCheck("blocking_check", () =>
+        {
+            return healthCheckTcs.Task;
+        });
+
+        var resource = builder.AddAzureCosmosDB("resource")
+                              .RunAsEmulator()
+                              .WithHealthCheck("blocking_check");
+
+        var dependentResource = builder.AddAzureCosmosDB("dependentresource")
+                                       .RunAsEmulator()
+                                       .WaitFor(resource);
+
+        using var app = builder.Build();
+
+        var pendingStart = app.StartAsync(cts.Token);
+
+        var rns = app.Services.GetRequiredService<ResourceNotificationService>();
+
+        await rns.WaitForResourceAsync(resource.Resource.Name, KnownResourceStates.Running, cts.Token);
+
+        await rns.WaitForResourceAsync(dependentResource.Resource.Name, KnownResourceStates.Waiting, cts.Token);
+
+        healthCheckTcs.SetResult(HealthCheckResult.Healthy());
+
+        await rns.WaitForResourceAsync(resource.Resource.Name, (re => re.Snapshot.HealthStatus == HealthStatus.Healthy), cts.Token);
+
+        await rns.WaitForResourceAsync(dependentResource.Resource.Name, KnownResourceStates.Running, cts.Token);
+
+        await pendingStart;
+
+        await app.StopAsync();
+    }
+
+}