File: AzureCosmosDBExtensions.cs
Web Access
Project: src\src\Aspire.Hosting.Azure.CosmosDB\Aspire.Hosting.Azure.CosmosDB.csproj (Aspire.Hosting.Azure.CosmosDB)
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Diagnostics.CodeAnalysis;
using System.Globalization;
using Aspire.Hosting.ApplicationModel;
using Aspire.Hosting.Azure;
using Aspire.Hosting.Azure.CosmosDB;
using Aspire.Hosting.Utils;
using Azure.Identity;
using Azure.Provisioning;
using Azure.Provisioning.CosmosDB;
using Azure.Provisioning.Expressions;
using Azure.Provisioning.KeyVault;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.DependencyInjection;
 
namespace Aspire.Hosting;
 
/// <summary>
/// Extension methods for adding Azure Cosmos DB resources to the application model.
/// </summary>
public static class AzureCosmosExtensions
{
    /// <summary>
    /// Adds an Azure Cosmos DB connection to the application model.
    /// </summary>
    /// <param name="builder">The <see cref="IDistributedApplicationBuilder"/>.</param>
    /// <param name="name">The name of the resource. This name will be used as the connection string name when referenced in a dependency.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
    public static IResourceBuilder<AzureCosmosDBResource> AddAzureCosmosDB(this IDistributedApplicationBuilder builder, [ResourceName] string name)
    {
        builder.AddAzureProvisioning();
 
        var resource = new AzureCosmosDBResource(name, ConfigureCosmosDBInfrastructure);
        return builder.AddResource(resource)
                      .WithManifestPublishingCallback(resource.WriteToManifest);
    }
 
    /// <summary>
    /// Configures an Azure Cosmos DB resource to be emulated using the Azure Cosmos DB emulator with the NoSQL API. This resource requires an <see cref="AzureCosmosDBResource"/> to be added to the application model.
    /// For more information on the Azure Cosmos DB emulator, see <a href="https://learn.microsoft.com/azure/cosmos-db/emulator#authentication"></a>.
    /// </summary>
    /// <param name="builder">The Azure Cosmos DB resource builder.</param>
    /// <param name="configureContainer">Callback that exposes underlying container used for emulation to allow for customization.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
    /// <remarks>
    /// When using the Azure Cosmos DB emulator, the container requires a TLS/SSL certificate.
    /// For more information, see <a href="https://learn.microsoft.com/azure/cosmos-db/how-to-develop-emulator?tabs=docker-linux#export-the-emulators-tlsssl-certificate"></a>.
    /// This version of the package defaults to the <inheritdoc cref="CosmosDBEmulatorContainerImageTags.Tag"/> tag of the <inheritdoc cref="CosmosDBEmulatorContainerImageTags.Registry"/>/<inheritdoc cref="CosmosDBEmulatorContainerImageTags.Image"/> container image.
    /// </remarks>
    public static IResourceBuilder<AzureCosmosDBResource> RunAsEmulator(this IResourceBuilder<AzureCosmosDBResource> builder, Action<IResourceBuilder<AzureCosmosDBEmulatorResource>>? configureContainer = null)
        => builder.RunAsEmulator(configureContainer, useVNextPreview: false);
 
    /// <summary>
    /// Configures an Azure Cosmos DB resource to be emulated using the Azure Cosmos DB Linux-based emulator (preview) with the NoSQL API. This resource requires an <see cref="AzureCosmosDBResource"/> to be added to the application model.
    /// For more information on the Azure Cosmos DB emulator, see <a href="https://learn.microsoft.com/azure/cosmos-db/emulator-linux"></a>.
    /// </summary>
    /// <param name="builder">The Azure Cosmos DB resource builder.</param>
    /// <param name="configureContainer">Callback that exposes underlying container used for emulation to allow for customization.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
    /// <remarks>
    /// This version of the package defaults to the <inheritdoc cref="CosmosDBEmulatorContainerImageTags.TagVNextPreview"/> tag of the <inheritdoc cref="CosmosDBEmulatorContainerImageTags.Registry"/>/<inheritdoc cref="CosmosDBEmulatorContainerImageTags.Image"/> container image.
    /// </remarks>
    [Experimental("ASPIRECOSMOS001", UrlFormat = "https://aka.ms/dotnet/aspire/diagnostics#{0}")]
    public static IResourceBuilder<AzureCosmosDBResource> RunAsPreviewEmulator(this IResourceBuilder<AzureCosmosDBResource> builder, Action<IResourceBuilder<AzureCosmosDBEmulatorResource>>? configureContainer = null)
        => builder.RunAsEmulator(configureContainer, useVNextPreview: true);
 
    private static IResourceBuilder<AzureCosmosDBResource> RunAsEmulator(this IResourceBuilder<AzureCosmosDBResource> builder, Action<IResourceBuilder<AzureCosmosDBEmulatorResource>>? configureContainer, bool useVNextPreview)
    {
        if (builder.ApplicationBuilder.ExecutionContext.IsPublishMode)
        {
            return builder;
        }
 
        var scheme = useVNextPreview ? "http" : null;
        builder.WithEndpoint(name: "emulator", scheme: scheme, targetPort: 8081)
               .WithAnnotation(new ContainerImageAnnotation
               {
                   Registry = CosmosDBEmulatorContainerImageTags.Registry,
                   Image = CosmosDBEmulatorContainerImageTags.Image,
                   Tag = useVNextPreview ? CosmosDBEmulatorContainerImageTags.TagVNextPreview : CosmosDBEmulatorContainerImageTags.Tag
               });
 
        CosmosClient? cosmosClient = null;
 
        builder.ApplicationBuilder.Eventing.Subscribe<ConnectionStringAvailableEvent>(builder.Resource, async (@event, ct) =>
        {
            var connectionString = await builder.Resource.ConnectionStringExpression.GetValueAsync(ct).ConfigureAwait(false);
 
            if (connectionString == null)
            {
                throw new DistributedApplicationException($"ConnectionStringAvailableEvent was published for the '{builder.Resource.Name}' resource but the connection string was null.");
            }
 
            cosmosClient = CreateCosmosClient(connectionString);
        });
 
        builder.ApplicationBuilder.Eventing.Subscribe<ResourceReadyEvent>(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().AddAzureCosmosDB(
            sp => cosmosClient ?? throw new InvalidOperationException("CosmosClient is not initialized."),
            name: healthCheckKey
            );
 
        builder.WithHealthCheck(healthCheckKey);
 
        if (configureContainer != null)
        {
            var surrogate = new AzureCosmosDBEmulatorResource(builder.Resource);
            var surrogateBuilder = builder.ApplicationBuilder.CreateResourceBuilder(surrogate);
            configureContainer(surrogateBuilder);
        }
 
        return builder;
 
        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>
    /// Adds a named volume for the data folder to an Azure Cosmos DB emulator resource.
    /// </summary>
    /// <param name="builder">The builder for the <see cref="AzureCosmosDBEmulatorResource"/>.</param>
    /// <param name="name">The name of the volume. Defaults to an auto-generated name based on the application and resource names.</param>
    /// <returns>A builder for the <see cref="AzureCosmosDBEmulatorResource"/>.</returns>
    public static IResourceBuilder<AzureCosmosDBEmulatorResource> WithDataVolume(this IResourceBuilder<AzureCosmosDBEmulatorResource> builder, string? name = null)
    {
        var dataPath = builder.Resource.InnerResource.IsPreviewEmulator ? "/data" : "/tmp/cosmos/appdata";
 
        return builder.WithEnvironment("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "true")
            .WithVolume(name ?? VolumeNameGenerator.Generate(builder, "data"), dataPath, isReadOnly: false);
    }
 
    /// <summary>
    /// Configures the gateway port for the Azure Cosmos DB emulator.
    /// </summary>
    /// <param name="builder">Builder for the Cosmos emulator container</param>
    /// <param name="port">Host port to bind to the emulator gateway port.</param>
    /// <returns>Cosmos emulator resource builder.</returns>
    public static IResourceBuilder<AzureCosmosDBEmulatorResource> WithGatewayPort(this IResourceBuilder<AzureCosmosDBEmulatorResource> builder, int? port)
    {
        return builder.WithEndpoint("emulator", endpoint =>
        {
            endpoint.Port = port;
        });
    }
 
    /// <summary>
    /// Configures the partition count for the Azure Cosmos DB emulator.
    /// </summary>
    /// <param name="builder">Builder for the Cosmos emulator container</param>
    /// <param name="count">Desired partition count.</param>
    /// <returns>Cosmos emulator resource builder.</returns>
    /// <remarks>Not calling this method will result in the default of 10 partitions. The actual started partitions is always one more than specified.
    /// See <a href="https://learn.microsoft.com/en-us/azure/cosmos-db/emulator-windows-arguments#change-the-number-of-default-containers">this documentation</a> about setting the partition count.
    /// </remarks>
    public static IResourceBuilder<AzureCosmosDBEmulatorResource> WithPartitionCount(this IResourceBuilder<AzureCosmosDBEmulatorResource> builder, int count)
    {
        if (builder.Resource.InnerResource.IsPreviewEmulator)
        {
            throw new NotSupportedException($"'{nameof(WithPartitionCount)}' does not work when using the preview version of the Azure Cosmos DB emulator.");
        }
 
        if (count < 1 || count > 250)
        {
            throw new ArgumentOutOfRangeException(nameof(count), count, "Count must be between 1 and 250.");
        }
 
        return builder.WithEnvironment("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", count.ToString(CultureInfo.InvariantCulture));
    }
 
    /// <summary>
    /// Adds a database to the associated Cosmos DB account resource.
    /// </summary>
    /// <param name="builder">AzureCosmosDB resource builder.</param>
    /// <param name="databaseName">Name of database.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
    [Obsolete($"This method is obsolete and will be removed in a future version. Use {nameof(WithDatabase)} instead to add a Cosmos DB database.")]
    public static IResourceBuilder<AzureCosmosDBResource> AddDatabase(this IResourceBuilder<AzureCosmosDBResource> builder, string databaseName)
    {
        return builder.WithDatabase(databaseName);
    }
 
    /// <summary>
    /// Adds a database to the associated Cosmos DB account resource.
    /// </summary>
    /// <param name="builder">AzureCosmosDB resource builder.</param>
    /// <param name="name">Name of database.</param>
    /// <param name="configure">An optional method that can be used for customizing the <see cref="CosmosDBDatabase"/>.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
    public static IResourceBuilder<AzureCosmosDBResource> WithDatabase(this IResourceBuilder<AzureCosmosDBResource> builder, string name, Action<CosmosDBDatabase>? configure = null)
    {
        var database = builder.Resource.Databases.FirstOrDefault(x => x.Name == name);
 
        if (database == null)
        {
            database = new CosmosDBDatabase(name);
            builder.Resource.Databases.Add(database);
        }
 
        configure?.Invoke(database);
        return builder;
    }
 
    /// <summary>
    /// Configures the Azure Cosmos DB preview emulator to expose the Data Explorer endpoint.
    /// </summary>
    /// <param name="builder">Builder for the Cosmos emulator container</param>
    /// <param name="port">Optional host port to bind the Data Explorer to.</param>
    /// <returns>Cosmos emulator resource builder.</returns>
    /// <remarks>
    /// The Data Explorer is only available with <see cref="RunAsPreviewEmulator"/>.
    /// </remarks>
    [Experimental("ASPIRECOSMOS001", UrlFormat = "https://aka.ms/dotnet/aspire/diagnostics#{0}")]
    public static IResourceBuilder<AzureCosmosDBEmulatorResource> WithDataExplorer(this IResourceBuilder<AzureCosmosDBEmulatorResource> builder, int? port = null)
    {
        if (!builder.Resource.InnerResource.IsPreviewEmulator)
        {
            throw new NotSupportedException($"The Data Explorer endpoint is only available when using the preview version of the Azure Cosmos DB emulator. Call '{nameof(RunAsPreviewEmulator)}' instead.");
        }
 
        return builder.WithEndpoint(endpointName: "data-explorer", endpoint =>
        {
            endpoint.UriScheme = "http";
            endpoint.TargetPort = 1234;
            endpoint.Port = port;
        });
    }
 
    /// <summary>    
    /// Configures the resource to use access key authentication with Azure Cosmos DB.
    /// </summary>
    /// <param name="builder">The Azure Cosmos DB resource builder.</param>
    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/> builder.</returns>
    /// <example>
    /// The following example creates an Azure Cosmos DB resource that uses access key authentication.
    /// <code lang="csharp">
    /// var builder = DistributedApplication.CreateBuilder(args);
    ///
    /// var cosmosdb = builder.AddAzureCosmosDB("cache")
    ///     .WithAccessKeyAuthentication();
    ///
    /// builder.AddProject&lt;Projects.ProductService&gt;()
    ///     .WithReference(cosmosdb);
    ///
    /// builder.Build().Run();
    /// </code>
    /// </example>
    public static IResourceBuilder<AzureCosmosDBResource> WithAccessKeyAuthentication(
        this IResourceBuilder<AzureCosmosDBResource> builder)
    {
        ArgumentNullException.ThrowIfNull(builder);
 
        var azureResource = builder.Resource;
        azureResource.ConnectionStringSecretOutput = new BicepSecretOutputReference("connectionString", azureResource);
 
        return builder;
    }
 
    private static void ConfigureCosmosDBInfrastructure(AzureResourceInfrastructure infrastructure)
    {
        var azureResource = (AzureCosmosDBResource)infrastructure.AspireResource;
 
        var cosmosAccount = new CosmosDBAccount(infrastructure.AspireResource.GetBicepIdentifier())
        {
            Kind = CosmosDBAccountKind.GlobalDocumentDB,
            ConsistencyPolicy = new ConsistencyPolicy()
            {
                DefaultConsistencyLevel = DefaultConsistencyLevel.Session
            },
            DatabaseAccountOfferType = CosmosDBAccountOfferType.Standard,
            Locations =
                {
                    new CosmosDBAccountLocation
                    {
                        LocationName = new IdentifierExpression("location"),
                        FailoverPriority = 0
                    }
                },
            Tags = { { "aspire-resource-name", infrastructure.AspireResource.Name } }
        };
        infrastructure.Add(cosmosAccount);
 
        foreach (var database in azureResource.Databases)
        {
            var cosmosSqlDatabase = new CosmosDBSqlDatabase(Infrastructure.NormalizeBicepIdentifier(database.Name))
            {
                Parent = cosmosAccount,
                Name = database.Name,
                Resource = new CosmosDBSqlDatabaseResourceInfo()
                {
                    DatabaseName = database.Name
                }
            };
            infrastructure.Add(cosmosSqlDatabase);
 
            foreach (var container in database.Containers)
            {
                var cosmosContainer = new CosmosDBSqlContainer(Infrastructure.NormalizeBicepIdentifier(container.Name))
                {
                    Parent = cosmosSqlDatabase,
                    Name = container.Name,
                    Resource = new CosmosDBSqlContainerResourceInfo()
                    {
                        ContainerName = container.Name,
                        PartitionKey = new CosmosDBContainerPartitionKey { Paths = [container.PartitionKeyPath] }
                    }
                };
                infrastructure.Add(cosmosContainer);
            }
        }
 
        if (azureResource.UseAccessKeyAuthentication)
        {
            cosmosAccount.DisableLocalAuth = false;
 
            var kvNameParam = new ProvisioningParameter(AzureBicepResource.KnownParameters.KeyVaultName, typeof(string));
            infrastructure.Add(kvNameParam);
 
            var keyVault = KeyVaultService.FromExisting("keyVault");
            keyVault.Name = kvNameParam;
            infrastructure.Add(keyVault);
 
            var secret = new KeyVaultSecret("connectionString")
            {
                Parent = keyVault,
                Name = "connectionString",
                Properties = new SecretProperties
                {
                    Value = BicepFunction.Interpolate($"AccountEndpoint={cosmosAccount.DocumentEndpoint};AccountKey={cosmosAccount.GetKeys().PrimaryMasterKey}")
                }
            };
            infrastructure.Add(secret);
        }
        else
        {
            cosmosAccount.DisableLocalAuth = true;
 
            var principalTypeParameter = new ProvisioningParameter(AzureBicepResource.KnownParameters.PrincipalType, typeof(string));
            infrastructure.Add(principalTypeParameter);
 
            var principalIdParameter = new ProvisioningParameter(AzureBicepResource.KnownParameters.PrincipalId, typeof(string));
            infrastructure.Add(principalIdParameter);
 
            var roleDefinition = CosmosDBSqlRoleDefinition_Derived.FromExisting(cosmosAccount.BicepIdentifier + "_roleDefinition");
            roleDefinition.Parent = cosmosAccount;
            roleDefinition.NameOverride = "00000000-0000-0000-0000-000000000002"; // data plane contributor role
            infrastructure.Add(roleDefinition);
 
            infrastructure.Add(new CosmosDBSqlRoleAssignment_Derived(cosmosAccount.BicepIdentifier + "_roleAssignment")
            {
                NameOverride = BicepFunction.CreateGuid(principalIdParameter, roleDefinition.Id, cosmosAccount.Id),
                Parent = cosmosAccount,
                Scope = cosmosAccount.Id,
                RoleDefinitionId = roleDefinition.Id,
                PrincipalId = principalIdParameter
            });
 
            infrastructure.Add(new ProvisioningOutput("connectionString", typeof(string))
            {
                Value = cosmosAccount.DocumentEndpoint
            });
        }
    }
}
 
// The following classes are working around https://github.com/Azure/azure-sdk-for-net/issues/47979 and can be removed once the issue is fixed.
 
internal class CosmosDBSqlRoleDefinition_Derived : CosmosDBSqlRoleDefinition
{
    private BicepValue<string>? _nameOverride;
 
    public CosmosDBSqlRoleDefinition_Derived(string name) : base(name)
    {
    }
 
    public static CosmosDBSqlRoleDefinition_Derived FromExisting(string bicepIdentifier)
    {
        return new CosmosDBSqlRoleDefinition_Derived(bicepIdentifier)
        {
            IsExistingResource = true
        };
    }
 
    public BicepValue<string> NameOverride
    {
        get
        {
            Initialize();
            return _nameOverride!;
        }
        set
        {
            Initialize();
            _nameOverride!.Assign(value);
        }
    }
 
    protected override void DefineProvisionableProperties()
    {
        base.DefineProvisionableProperties();
 
        _nameOverride = DefineProperty<string>("Name", new string[1] { "name" });
    }
}
 
internal class CosmosDBSqlRoleAssignment_Derived : CosmosDBSqlRoleAssignment
{
    private BicepValue<string>? _nameOverride;
 
    public CosmosDBSqlRoleAssignment_Derived(string name) : base(name)
    {
    }
 
    public BicepValue<string> NameOverride
    {
        get
        {
            Initialize();
            return _nameOverride!;
        }
        set
        {
            Initialize();
            _nameOverride!.Assign(value);
        }
    }
 
    protected override void DefineProvisionableProperties()
    {
        base.DefineProvisionableProperties();
 
        _nameOverride = DefineProperty<string>("Name", new string[1] { "name" });
    }
}