From 2cb078e2829b1b5f0104f5d83b5a65e7c21d9070 Mon Sep 17 00:00:00 2001 From: Mitch Denny Date: Wed, 18 Sep 2024 05:38:08 +1000 Subject: [PATCH] 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 | 1 + .../CosmosEndToEnd.AppHost/Program.cs | 2 +- .../Aspire.Hosting.Azure.CosmosDB.csproj | 3 + .../AzureCosmosDBExtensions.cs | 49 ++++++++++++++- .../AzureCosmosDBEmulatorFunctionalTests.cs | 59 +++++++++++++++++++ 5 files changed, 112 insertions(+), 2 deletions(-) create mode 100644 tests/Aspire.Hosting.Azure.Tests/AzureCosmosDBEmulatorFunctionalTests.cs diff --git a/Directory.Packages.props b/Directory.Packages.props index 0cf339bcb9..1032128360 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -70,6 +70,7 @@ + diff --git a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs b/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs index 5271a2fea5..bf2cbff2c9 100644 --- a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs +++ b/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs @@ -9,7 +9,7 @@ builder.AddProject("api") .WithExternalHttpEndpoints() - .WithReference(db); + .WithReference(db).WaitFor(db); #if !SKIP_DASHBOARD_REFERENCE // This project is only added in playground projects to support development/debugging 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 8c225ab9bb..e0fa234173 100644 --- a/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj +++ b/src/Aspire.Hosting.Azure.CosmosDB/Aspire.Hosting.Azure.CosmosDB.csproj @@ -14,10 +14,13 @@ + + + diff --git a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs b/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs index 64fdf876c5..91fb4fe2cd 100644 --- a/src/Aspire.Hosting.Azure.CosmosDB/AzureCosmosDBExtensions.cs +++ b/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 AddAzureCosmosDB(this IDis }; var resource = new AzureCosmosDBResource(name, configureConstruct); + + CosmosClient? cosmosClient = null; + + builder.Eventing.Subscribe(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); + } + } } /// diff --git a/tests/Aspire.Hosting.Azure.Tests/AzureCosmosDBEmulatorFunctionalTests.cs b/tests/Aspire.Hosting.Azure.Tests/AzureCosmosDBEmulatorFunctionalTests.cs new file mode 100644 index 0000000000..9bd631a85c --- /dev/null +++ b/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(); + 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(); + + 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(); + } + +}