Add Cosmos database builder pattern for child containers (dotnet#8266)

Author: captainsafia

playground/CosmosEndToEnd/CosmosEndToEnd.ApiService/Program.cs

@@ -8,8 +8,9 @@
 var builder = WebApplication.CreateBuilder(args);
 
 builder.AddServiceDefaults();
-builder.AddAzureCosmosDatabase("db");
-builder.AddAzureCosmosContainer("entries");
+builder.AddAzureCosmosDatabase("db")
+    .AddKeyedContainer("entries")
+    .AddKeyedContainer("users");
 builder.AddCosmosDbContext<TestCosmosContext>("db", configureDbContextOptions =>
 {
     configureDbContextOptions.RequestTimeout = TimeSpan.FromSeconds(120);
@@ -18,14 +19,13 @@
 var app = builder.Build();
 
 app.MapDefaultEndpoints();
-app.MapGet("/", async (Database db, Container container) =>
+
+static async Task<object> AddAndGetStatus<T>(Container container, T newEntry)
 {
-    // Add an entry to the database on each request.
-    var newEntry = new Entry() { Id = Guid.NewGuid().ToString() };
     await container.CreateItemAsync(newEntry);
 
-    var entries = new List<Entry>();
-    var iterator = container.GetItemQueryIterator<Entry>(requestOptions: new QueryRequestOptions() { MaxItemCount = 5 });
+    var entries = new List<T>();
+    var iterator = container.GetItemQueryIterator<T>(requestOptions: new QueryRequestOptions() { MaxItemCount = 5 });
 
     var batchCount = 0;
     while (iterator.HasMoreResults)
@@ -40,10 +40,22 @@
 
     return new
     {
-        batchCount = batchCount,
+        batchCount,
         totalEntries = entries.Count,
-        entries = entries
+        entries
     };
+}
+
+app.MapGet("/", async ([FromKeyedServices("entries")] Container container) =>
+{
+    var newEntry = new Entry() { Id = Guid.NewGuid().ToString() };
+    return await AddAndGetStatus(container, newEntry);
+});
+
+app.MapGet("/users", async ([FromKeyedServices("users")] Container container) =>
+{
+    var newEntry = new User() { Id = $"user-{Guid.NewGuid()}" };
+    return await AddAndGetStatus(container, newEntry);
 });
 
 app.MapGet("/ef", async (TestCosmosContext context) =>
@@ -58,6 +70,12 @@
 
 app.Run();
 
+public class User
+{
+    [JsonProperty("id")]
+    public string? Id { get; set; }
+}
+
 public class Entry
 {
     [JsonProperty("id")]

playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs

@@ -9,12 +9,13 @@
                 .RunAsPreviewEmulator(e => e.WithDataExplorer());
 
 var db = cosmos.AddCosmosDatabase("db");
-var container = db.AddContainer("entries", "/id");
+var entries = db.AddContainer("entries", "/id", "staging-entries");
+db.AddContainer("users", "/id");
 
 builder.AddProject<Projects.CosmosEndToEnd_ApiService>("api")
        .WithExternalHttpEndpoints()
        .WithReference(db).WaitFor(db)
-       .WithReference(container).WaitFor(container);
+       .WithReference(entries).WaitFor(entries);
 
 #if !SKIP_DASHBOARD_REFERENCE
 // This project is only added in playground projects to support development/debugging

src/Components/Aspire.Microsoft.Azure.Cosmos/AspireMicrosoftAzureCosmosExtensions.cs

@@ -38,39 +38,6 @@ public static void AddAzureCosmosClient(
         builder.Services.AddSingleton(sp => GetCosmosClient(connectionName, settings, clientOptions));
     }
 
-    /// <summary>
-    /// Registers the <see cref="Database"/> as a singleton in the services provided by the <paramref name="builder"/>.
-    /// </summary>
-    /// <param name="builder">The <see cref="IHostApplicationBuilder" /> to read config from and add services to.</param>
-    /// <param name="connectionName">The connection name to use to find a connection string.</param>
-    /// <param name="configureSettings">An optional method that can be used for customizing the <see cref="MicrosoftAzureCosmosSettings"/>. It's invoked after the settings are read from the configuration.</param>
-    /// <param name="configureClientOptions">An optional method that can be used for customizing the <see cref="CosmosClientOptions"/>.</param>
-    /// <remarks>Reads the configuration from "Aspire:Microsoft:Azure:Cosmos" section.</remarks>
-    /// <exception cref="InvalidOperationException">If required ConnectionString is not provided in configuration section</exception>
-    public static void AddAzureCosmosDatabase(
-        this IHostApplicationBuilder builder,
-        string connectionName,
-        Action<MicrosoftAzureCosmosSettings>? configureSettings = null,
-        Action<CosmosClientOptions>? configureClientOptions = null)
-    {
-        var settings = builder.GetSettings(connectionName, configureSettings);
-        var clientOptions = builder.GetClientOptions(settings, configureClientOptions);
-        builder.Services.AddSingleton(sp =>
-        {
-            if (string.IsNullOrEmpty(settings.DatabaseName))
-            {
-                throw new InvalidOperationException($"The connection string '{connectionName}' does not exist or is missing the database name.");
-            }
-            CosmosClient? client = null;
-            if (configureClientOptions is null)
-            {
-                client = sp.GetService<CosmosClient>();
-            }
-            client ??= GetCosmosClient(connectionName, settings, clientOptions);
-            return client.GetDatabase(settings.DatabaseName);
-        });
-    }
-
     /// <summary>
     /// Registers the <see cref="Container"/> as a singleton in the services provided by the <paramref name="builder"/>.
     /// </summary>
@@ -79,6 +46,15 @@ public static void AddAzureCosmosDatabase(
     /// <param name="configureSettings">An optional method that can be used for customizing the <see cref="MicrosoftAzureCosmosSettings"/>. It's invoked after the settings are read from the configuration.</param>
     /// <param name="configureClientOptions">An optional method that can be used for customizing the <see cref="CosmosClientOptions"/>.</param>
     /// <remarks>Reads the configuration from "Aspire:Microsoft:Azure:Cosmos" section.</remarks>
+    /// <remarks>
+    /// The <see cref="Container"/> is registered as a singleton in the services provided by
+    /// the <paramref name="builder"/> and does not reuse any existing <see cref="CosmosClient"/>
+    /// instances in the DI container. The connection string associated with the <paramref name="connectionName"/>
+    /// must contain the database name and container name or be set in the <paramref name="configureSettings" />
+    /// callback. To interact with multiple containers against the same database, use
+    /// <see cref="CosmosDatabaseBuilder"/> to register the database and then call
+    /// <see cref="CosmosDatabaseBuilder.AddKeyedContainer(string)"/> for each container.
+    /// </remarks>
     /// <exception cref="InvalidOperationException">If required ConnectionString is not provided in configuration section</exception>
     public static void AddAzureCosmosContainer(
         this IHostApplicationBuilder builder,
@@ -94,12 +70,7 @@ public static void AddAzureCosmosContainer(
             {
                 throw new InvalidOperationException($"The connection string '{connectionName}' does not exist or is missing the container name or database name.");
             }
-            CosmosClient? client = null;
-            if (configureClientOptions is null)
-            {
-                client = sp.GetService<CosmosClient>();
-            }
-            client ??= GetCosmosClient(connectionName, settings, clientOptions);
+            var client = GetCosmosClient(connectionName, settings, clientOptions);
             return client.GetContainer(settings.DatabaseName, settings.ContainerName);
         });
     }
@@ -132,72 +103,89 @@ public static void AddKeyedAzureCosmosClient(
     }
 
     /// <summary>
-    /// Registers the <see cref="Database"/> as a singleton for given <paramref name="name" /> in the services provided by the <paramref name="builder"/>.
+    /// Registers the <see cref="Container"/> as a singleton for given <paramref name="name" /> in the services provided by the <paramref name="builder"/>.
     /// </summary>
     /// <param name="builder">The <see cref="IHostApplicationBuilder" /> to read config from and add services to.</param>
     /// <param name="name">The name of the component, which is used as the <see cref="ServiceDescriptor.ServiceKey"/> of the service and also to retrieve the connection string from the ConnectionStrings configuration section.</param>
     /// <param name="configureSettings">An optional method that can be used for customizing the <see cref="MicrosoftAzureCosmosSettings"/>. It's invoked after the settings are read from the configuration.</param>
     /// <param name="configureClientOptions">An optional method that can be used for customizing the <see cref="CosmosClientOptions"/>.</param>
     /// <remarks>Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section.</remarks>
+    /// <remarks>
+    /// The <see cref="Container"/> is registered as a singleton in the services provided by
+    /// the <paramref name="builder"/> and does not reuse any existing <see cref="CosmosClient"/>
+    /// instances in the DI container. The connection string associated with the <paramref name="name"/>
+    /// must contain the database name and container name or be set in the <paramref name="configureSettings" />
+    /// callback. To interact with multiple containers against the same database, use
+    /// <see cref="CosmosDatabaseBuilder"/> to register the database and then call
+    /// <see cref="CosmosDatabaseBuilder.AddKeyedContainer(string)"/> for each container.
+    /// </remarks>
     /// <exception cref="InvalidOperationException">If required ConnectionString is not provided in configuration section</exception>
-    public static void AddKeyedAzureCosmosDatabase(
-       this IHostApplicationBuilder builder,
-       string name,
-       Action<MicrosoftAzureCosmosSettings>? configureSettings = null,
-       Action<CosmosClientOptions>? configureClientOptions = null)
+    public static void AddKeyedAzureCosmosContainer(
+        this IHostApplicationBuilder builder,
+        string name,
+        Action<MicrosoftAzureCosmosSettings>? configureSettings = null,
+        Action<CosmosClientOptions>? configureClientOptions = null)
     {
         var settings = builder.GetSettings(name, configureSettings);
         var clientOptions = builder.GetClientOptions(settings, configureClientOptions);
         builder.Services.AddKeyedSingleton(name, (sp, key) =>
         {
-            if (string.IsNullOrEmpty(settings.DatabaseName))
-            {
-                throw new InvalidOperationException($"The connection string '{name}' does not exist or is missing the database name.");
-            }
-            CosmosClient? client = null;
-            if (configureClientOptions is null)
+            if (string.IsNullOrEmpty(settings.ContainerName) || string.IsNullOrEmpty(settings.DatabaseName))
             {
-                client = sp.GetKeyedService<CosmosClient>(key);
+                throw new InvalidOperationException($"The connection string '{name}' does not exist or is missing the container name or database name.");
             }
-            client ??= GetCosmosClient(name, settings, clientOptions);
-            return client.GetDatabase(settings.DatabaseName);
+            var client = GetCosmosClient(name, settings, clientOptions);
+            return client.GetContainer(settings.DatabaseName, settings.ContainerName);
         });
     }
 
     /// <summary>
-    /// Registers the <see cref="Container"/> as a singleton for given <paramref name="name" /> in the services provided by the <paramref name="builder"/>.
+    /// Registers the <see cref="Database"/> as a singleton the services provided by the <paramref name="builder"/>
+    /// and returns a <see cref="CosmosDatabaseBuilder"/> to support chaining multiple container registrations against the same database.
     /// </summary>
     /// <param name="builder">The <see cref="IHostApplicationBuilder" /> to read config from and add services to.</param>
-    /// <param name="name">The name of the component, which is used as the <see cref="ServiceDescriptor.ServiceKey"/> of the service and also to retrieve the connection string from the ConnectionStrings configuration section.</param>
+    /// <param name="connectionName">The connection name to use to find a connection string.</param>
     /// <param name="configureSettings">An optional method that can be used for customizing the <see cref="MicrosoftAzureCosmosSettings"/>. It's invoked after the settings are read from the configuration.</param>
     /// <param name="configureClientOptions">An optional method that can be used for customizing the <see cref="CosmosClientOptions"/>.</param>
     /// <remarks>Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section.</remarks>
     /// <exception cref="InvalidOperationException">If required ConnectionString is not provided in configuration section</exception>
-    public static void AddKeyedAzureCosmosContainer(
+    public static CosmosDatabaseBuilder AddAzureCosmosDatabase(
         this IHostApplicationBuilder builder,
-        string name,
+        string connectionName,
         Action<MicrosoftAzureCosmosSettings>? configureSettings = null,
         Action<CosmosClientOptions>? configureClientOptions = null)
+    {
+        var settings = builder.GetSettings(connectionName, configureSettings);
+        var clientOptions = builder.GetClientOptions(settings, configureClientOptions);
+        var cosmosDatabaseBuilder = new CosmosDatabaseBuilder(builder, connectionName, settings, clientOptions);
+        cosmosDatabaseBuilder.AddDatabase();
+        return cosmosDatabaseBuilder;
+    }
+
+    /// <summary>
+    /// Registers the <see cref="Database"/> as a singleton for given <paramref name="name" /> in the services provided by the <paramref name="builder"/>
+    /// and returns a <see cref="CosmosDatabaseBuilder"/> to support chaining multiple container registrations against the same database.
+    /// </summary>
+    /// <param name="builder">The <see cref="IHostApplicationBuilder" /> to read config from and add services to.</param>
+    /// <param name="name">The name of the component, which is used as the <see cref="ServiceDescriptor.ServiceKey"/> of the service and also to retrieve the connection string from the ConnectionStrings configuration section.</param>
+    /// <param name="configureSettings">An optional method that can be used for customizing the <see cref="MicrosoftAzureCosmosSettings"/>. It's invoked after the settings are read from the configuration.</param>
+    /// <param name="configureClientOptions">An optional method that can be used for customizing the <see cref="CosmosClientOptions"/>.</param>
+    /// <remarks>Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section.</remarks>
+    /// <exception cref="InvalidOperationException">If required ConnectionString is not provided in configuration section</exception>
+    public static CosmosDatabaseBuilder AddKeyedAzureCosmosDatabase(
+       this IHostApplicationBuilder builder,
+       string name,
+       Action<MicrosoftAzureCosmosSettings>? configureSettings = null,
+       Action<CosmosClientOptions>? configureClientOptions = null)
     {
         var settings = builder.GetSettings(name, configureSettings);
         var clientOptions = builder.GetClientOptions(settings, configureClientOptions);
-        builder.Services.AddKeyedSingleton(name, (sp, key) =>
-        {
-            if (string.IsNullOrEmpty(settings.ContainerName) || string.IsNullOrEmpty(settings.DatabaseName))
-            {
-                throw new InvalidOperationException($"The connection string '{name}' does not exist or is missing the container name or database name.");
-            }
-            CosmosClient? client = null;
-            if (configureClientOptions is null)
-            {
-                client = sp.GetKeyedService<CosmosClient>(key);
-            }
-            client ??= GetCosmosClient(name, settings, clientOptions);
-            return client.GetContainer(settings.DatabaseName, settings.ContainerName);
-        });
+        var cosmosDatabaseBuilder = new CosmosDatabaseBuilder(builder, name, settings, clientOptions);
+        cosmosDatabaseBuilder.AddKeyedDatabase();
+        return cosmosDatabaseBuilder;
     }
 
-    private static CosmosConnectionInfo? GetCosmosConnectionInfo(this IHostApplicationBuilder builder, string connectionName)
+    internal static CosmosConnectionInfo? GetCosmosConnectionInfo(this IHostApplicationBuilder builder, string connectionName)
     {
         ArgumentNullException.ThrowIfNull(builder);
         ArgumentException.ThrowIfNullOrEmpty(connectionName);
@@ -275,7 +263,7 @@ private static CosmosClientOptions GetClientOptions(
         return clientOptions;
     }
 
-    private static CosmosClient GetCosmosClient(string connectionName, MicrosoftAzureCosmosSettings settings, CosmosClientOptions clientOptions)
+    internal static CosmosClient GetCosmosClient(string connectionName, MicrosoftAzureCosmosSettings settings, CosmosClientOptions clientOptions)
     {
         if (!string.IsNullOrEmpty(settings.ConnectionString))
         {

src/Components/Aspire.Microsoft.Azure.Cosmos/CosmosDatabaseBuilder.cs

@@ -0,0 +1,81 @@
+// 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.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+
+namespace Aspire.Microsoft.Azure.Cosmos;
+
+/// <summary>
+/// Represents a builder that can be used to register multiple container
+/// instances against the same Cosmos database connection.
+/// </summary>
+public sealed class CosmosDatabaseBuilder(
+    IHostApplicationBuilder hostBuilder,
+    string connectionName,
+    MicrosoftAzureCosmosSettings settings,
+    CosmosClientOptions clientOptions)
+{
+    private CosmosClient? _client;
+
+    internal CosmosDatabaseBuilder AddDatabase()
+    {
+        hostBuilder.Services.AddSingleton(sp =>
+        {
+            if (string.IsNullOrEmpty(settings.DatabaseName))
+            {
+                throw new InvalidOperationException(
+                    $"A Database could not be configured. Ensure valid connection information was provided in 'ConnectionStrings:{connectionName}'.");
+            }
+            _client ??= AspireMicrosoftAzureCosmosExtensions.GetCosmosClient(connectionName, settings, clientOptions);
+            return _client.GetDatabase(settings.DatabaseName);
+        });
+
+        return this;
+    }
+
+    internal CosmosDatabaseBuilder AddKeyedDatabase()
+    {
+        hostBuilder.Services.AddKeyedSingleton(connectionName, (sp, _) =>
+        {
+            if (string.IsNullOrEmpty(settings.DatabaseName))
+            {
+                throw new InvalidOperationException(
+                    $"A Database could not be configured. Ensure valid connection information was provided in 'ConnectionStrings:{connectionName}'.");
+            }
+            _client ??= AspireMicrosoftAzureCosmosExtensions.GetCosmosClient(connectionName, settings, clientOptions);
+            return _client.GetDatabase(settings.DatabaseName);
+        });
+
+        return this;
+    }
+
+    /// <summary>
+    /// Register a <see cref="Container"/> against the database managed with <see cref="CosmosDatabaseBuilder"/> as a
+    /// keyed singleton.
+    /// </summary>
+    /// <param name="name">The name of the container to register.</param>
+    /// <returns>A <see cref="CosmosDatabaseBuilder"/> that can be used for further chaining.</returns>
+    public CosmosDatabaseBuilder AddKeyedContainer(string name)
+    {
+        _client ??= AspireMicrosoftAzureCosmosExtensions.GetCosmosClient(connectionName, settings, clientOptions);
+
+        var connectionInfo = hostBuilder.GetCosmosConnectionInfo(name);
+
+        hostBuilder.Services.AddKeyedSingleton(name, (sp, _) =>
+        {
+            // If a connection string was provided, check that it contains a valid container name.
+            if (connectionInfo is not null && string.IsNullOrEmpty(connectionInfo?.ContainerName))
+            {
+                throw new InvalidOperationException(
+                    $"A Container could not be configured. Ensure valid connection information was provided in 'ConnectionStrings:{name}'");
+            }
+
+            // Use the container name from the connection string if provided, otherwise use the name
+            return _client.GetContainer(settings.DatabaseName, connectionInfo?.ContainerName ?? name);
+        });
+
+        return this;
+    }
+}