diff --git a/playground/CosmosEndToEnd/CosmosEndToEnd.ApiService/Program.cs b/playground/CosmosEndToEnd/CosmosEndToEnd.ApiService/Program.cs index 47afc79fd04..5ec06da140a 100644 --- a/playground/CosmosEndToEnd/CosmosEndToEnd.ApiService/Program.cs +++ b/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("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 AddAndGetStatus(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(); - var iterator = container.GetItemQueryIterator(requestOptions: new QueryRequestOptions() { MaxItemCount = 5 }); + var entries = new List(); + var iterator = container.GetItemQueryIterator(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")] diff --git a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs b/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs index 97db04b0684..5a81fa6e9bc 100644 --- a/playground/CosmosEndToEnd/CosmosEndToEnd.AppHost/Program.cs +++ b/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("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 diff --git a/src/Components/Aspire.Microsoft.Azure.Cosmos/AspireMicrosoftAzureCosmosExtensions.cs b/src/Components/Aspire.Microsoft.Azure.Cosmos/AspireMicrosoftAzureCosmosExtensions.cs index af832a32495..61428f9f824 100644 --- a/src/Components/Aspire.Microsoft.Azure.Cosmos/AspireMicrosoftAzureCosmosExtensions.cs +++ b/src/Components/Aspire.Microsoft.Azure.Cosmos/AspireMicrosoftAzureCosmosExtensions.cs @@ -38,39 +38,6 @@ public static void AddAzureCosmosClient( builder.Services.AddSingleton(sp => GetCosmosClient(connectionName, settings, clientOptions)); } - /// - /// Registers the as a singleton in the services provided by the . - /// - /// The to read config from and add services to. - /// The connection name to use to find a connection string. - /// An optional method that can be used for customizing the . It's invoked after the settings are read from the configuration. - /// An optional method that can be used for customizing the . - /// Reads the configuration from "Aspire:Microsoft:Azure:Cosmos" section. - /// If required ConnectionString is not provided in configuration section - public static void AddAzureCosmosDatabase( - this IHostApplicationBuilder builder, - string connectionName, - Action? configureSettings = null, - Action? 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(); - } - client ??= GetCosmosClient(connectionName, settings, clientOptions); - return client.GetDatabase(settings.DatabaseName); - }); - } - /// /// Registers the as a singleton in the services provided by the . /// @@ -79,6 +46,15 @@ public static void AddAzureCosmosDatabase( /// An optional method that can be used for customizing the . It's invoked after the settings are read from the configuration. /// An optional method that can be used for customizing the . /// Reads the configuration from "Aspire:Microsoft:Azure:Cosmos" section. + /// + /// The is registered as a singleton in the services provided by + /// the and does not reuse any existing + /// instances in the DI container. The connection string associated with the + /// must contain the database name and container name or be set in the + /// callback. To interact with multiple containers against the same database, use + /// to register the database and then call + /// for each container. + /// /// If required ConnectionString is not provided in configuration section 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(); - } - 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( } /// - /// Registers the as a singleton for given in the services provided by the . + /// Registers the as a singleton for given in the services provided by the . /// /// The to read config from and add services to. /// The name of the component, which is used as the of the service and also to retrieve the connection string from the ConnectionStrings configuration section. /// An optional method that can be used for customizing the . It's invoked after the settings are read from the configuration. /// An optional method that can be used for customizing the . /// Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section. + /// + /// The is registered as a singleton in the services provided by + /// the and does not reuse any existing + /// instances in the DI container. The connection string associated with the + /// must contain the database name and container name or be set in the + /// callback. To interact with multiple containers against the same database, use + /// to register the database and then call + /// for each container. + /// /// If required ConnectionString is not provided in configuration section - public static void AddKeyedAzureCosmosDatabase( - this IHostApplicationBuilder builder, - string name, - Action? configureSettings = null, - Action? configureClientOptions = null) + public static void AddKeyedAzureCosmosContainer( + this IHostApplicationBuilder builder, + string name, + Action? configureSettings = null, + Action? 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(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); }); } /// - /// Registers the as a singleton for given in the services provided by the . + /// Registers the as a singleton the services provided by the + /// and returns a to support chaining multiple container registrations against the same database. /// /// The to read config from and add services to. - /// The name of the component, which is used as the of the service and also to retrieve the connection string from the ConnectionStrings configuration section. + /// The connection name to use to find a connection string. /// An optional method that can be used for customizing the . It's invoked after the settings are read from the configuration. /// An optional method that can be used for customizing the . /// Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section. /// If required ConnectionString is not provided in configuration section - public static void AddKeyedAzureCosmosContainer( + public static CosmosDatabaseBuilder AddAzureCosmosDatabase( this IHostApplicationBuilder builder, - string name, + string connectionName, Action? configureSettings = null, Action? 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; + } + + /// + /// Registers the as a singleton for given in the services provided by the + /// and returns a to support chaining multiple container registrations against the same database. + /// + /// The to read config from and add services to. + /// The name of the component, which is used as the of the service and also to retrieve the connection string from the ConnectionStrings configuration section. + /// An optional method that can be used for customizing the . It's invoked after the settings are read from the configuration. + /// An optional method that can be used for customizing the . + /// Reads the configuration from "Aspire:Microsoft:Azure:Cosmos:{name}" section. + /// If required ConnectionString is not provided in configuration section + public static CosmosDatabaseBuilder AddKeyedAzureCosmosDatabase( + this IHostApplicationBuilder builder, + string name, + Action? configureSettings = null, + Action? 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(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)) { diff --git a/src/Components/Aspire.Microsoft.Azure.Cosmos/CosmosDatabaseBuilder.cs b/src/Components/Aspire.Microsoft.Azure.Cosmos/CosmosDatabaseBuilder.cs new file mode 100644 index 00000000000..c1d012727ec --- /dev/null +++ b/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; + +/// +/// Represents a builder that can be used to register multiple container +/// instances against the same Cosmos database connection. +/// +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; + } + + /// + /// Register a against the database managed with as a + /// keyed singleton. + /// + /// The name of the container to register. + /// A that can be used for further chaining. + 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; + } +} diff --git a/tests/Aspire.Microsoft.Azure.Cosmos.Tests/AspireMicrosoftAzureCosmosExtensionsTests.cs b/tests/Aspire.Microsoft.Azure.Cosmos.Tests/AspireMicrosoftAzureCosmosExtensionsTests.cs index 8eff8388d12..8b09c8a0fc9 100644 --- a/tests/Aspire.Microsoft.Azure.Cosmos.Tests/AspireMicrosoftAzureCosmosExtensionsTests.cs +++ b/tests/Aspire.Microsoft.Azure.Cosmos.Tests/AspireMicrosoftAzureCosmosExtensionsTests.cs @@ -201,7 +201,7 @@ public void AddAzureCosmosDatabase_ThrowsForInvalidConnectionString(string conne using var host = builder.Build(); var exception = Assert.Throws(host.Services.GetRequiredService); - Assert.Equal("The connection string 'cosmos-key' does not exist or is missing the database name.", exception.Message); + Assert.Equal("A Database could not be configured. Ensure valid connection information was provided in 'ConnectionStrings:cosmos-key'.", exception.Message); } [Theory] @@ -217,7 +217,7 @@ public void AddKeyedAzureCosmosDatabase_ThrowsForInvalidConnectionString(string using var host = builder.Build(); var exception = Assert.Throws(() => host.Services.GetRequiredKeyedService(serviceKey)); - Assert.Contains("The connection string 'cosmos-key' does not exist or is missing the database name.", exception.Message); + Assert.Equal("A Database could not be configured. Ensure valid connection information was provided in 'ConnectionStrings:cosmos-key'.", exception.Message); } [Fact] @@ -242,12 +242,13 @@ public void AddAzureCosmosClient_RespectsLimitToEndpointViaConfigureSettings() } [Fact] - public void AddAzureCosmosDatabase_ReusesExistingClient() + public void AddAzureCosmosContainer_DoesNoReuseExistingClient() { var builder = Host.CreateEmptyApplicationBuilder(null); var databaseName = "testdb"; + var containerName = "testcontainer"; var expectedEndpoint = "https://localhost:8081/"; - var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};"; + var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};Container={containerName};"; PopulateConfiguration(builder.Configuration, connectionString); @@ -255,21 +256,45 @@ public void AddAzureCosmosDatabase_ReusesExistingClient() { options.LimitToEndpoint = false; }); - builder.AddAzureCosmosDatabase("cosmos"); + builder.AddAzureCosmosContainer("cosmos"); + + using var host = builder.Build(); + + var client = host.Services.GetRequiredService(); + var container = host.Services.GetRequiredService(); + + Assert.NotSame(client, container.Database.Client); + } + + [Fact] + public void AddAzureCosmosDatabase_CreatesNewClient_WhenConfigureClientOptionsProvided() + { + var builder = Host.CreateEmptyApplicationBuilder(null); + var databaseName = "testdb"; + var expectedEndpoint = "https://localhost:8081/"; + var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};"; + + PopulateConfiguration(builder.Configuration, connectionString); + + builder.AddAzureCosmosClient("cosmos"); + builder.AddAzureCosmosDatabase("cosmos", configureClientOptions: options => + { + options.LimitToEndpoint = false; + }); using var host = builder.Build(); var client = host.Services.GetRequiredService(); var database = host.Services.GetRequiredService(); - Assert.Same(client, database.Client); + Assert.NotSame(client, database.Client); Assert.Equal(databaseName, database.Id); Assert.Equal(expectedEndpoint, database.Client.Endpoint.ToString()); Assert.False(database.Client.ClientOptions.LimitToEndpoint); } [Fact] - public void AddAzureCosmosContainer_ReusesExistingClient() + public void AddAzureCosmosContainer_CreatesNewClient_WhenConfigureClientOptionsProvided() { var builder = Host.CreateEmptyApplicationBuilder(null); var databaseName = "testdb"; @@ -279,25 +304,25 @@ public void AddAzureCosmosContainer_ReusesExistingClient() PopulateConfiguration(builder.Configuration, connectionString); - builder.AddAzureCosmosClient("cosmos", configureClientOptions: options => + builder.AddAzureCosmosClient("cosmos"); + builder.AddAzureCosmosContainer("cosmos", configureClientOptions: options => { options.LimitToEndpoint = false; }); - builder.AddAzureCosmosContainer("cosmos"); using var host = builder.Build(); var client = host.Services.GetRequiredService(); var container = host.Services.GetRequiredService(); - Assert.Same(client, container.Database.Client); + Assert.NotSame(client, container.Database.Client); Assert.Equal(containerName, container.Id); Assert.Equal(expectedEndpoint, container.Database.Client.Endpoint.ToString()); Assert.False(container.Database.Client.ClientOptions.LimitToEndpoint); } [Fact] - public void AddKeyedAzureCosmosDatabase_ReusesExistingKeyedClient() + public void AddKeyedAzureCosmosDatabase_CreatesNewClient_WhenConfigureClientOptionsProvided() { var builder = Host.CreateEmptyApplicationBuilder(null); var serviceKey = "cosmos-key"; @@ -307,25 +332,25 @@ public void AddKeyedAzureCosmosDatabase_ReusesExistingKeyedClient() PopulateConfiguration(builder.Configuration, connectionString, serviceKey); - builder.AddKeyedAzureCosmosClient(serviceKey, configureClientOptions: options => + builder.AddKeyedAzureCosmosClient(serviceKey); + builder.AddKeyedAzureCosmosDatabase(serviceKey, configureClientOptions: options => { options.LimitToEndpoint = false; }); - builder.AddKeyedAzureCosmosDatabase(serviceKey); using var host = builder.Build(); var client = host.Services.GetRequiredKeyedService(serviceKey); var database = host.Services.GetRequiredKeyedService(serviceKey); - Assert.Same(client, database.Client); + Assert.NotSame(client, database.Client); Assert.Equal(databaseName, database.Id); Assert.Equal(expectedEndpoint, database.Client.Endpoint.ToString()); Assert.False(database.Client.ClientOptions.LimitToEndpoint); } [Fact] - public void AddKeyedAzureCosmosContainer_ReusesExistingKeyedClient() + public void AddKeyedAzureCosmosContainer_CreatesNewClient_WhenConfigureClientOptionsProvided() { var builder = Host.CreateEmptyApplicationBuilder(null); var serviceKey = "cosmos-key"; @@ -336,24 +361,25 @@ public void AddKeyedAzureCosmosContainer_ReusesExistingKeyedClient() PopulateConfiguration(builder.Configuration, connectionString, serviceKey); - builder.AddKeyedAzureCosmosClient(serviceKey, configureClientOptions: options => + builder.AddKeyedAzureCosmosClient(serviceKey); + builder.AddKeyedAzureCosmosContainer(serviceKey, configureClientOptions: options => { options.LimitToEndpoint = false; }); - builder.AddKeyedAzureCosmosContainer(serviceKey); using var host = builder.Build(); var client = host.Services.GetRequiredKeyedService(serviceKey); var container = host.Services.GetRequiredKeyedService(serviceKey); - Assert.Same(client, container.Database.Client); + Assert.NotSame(client, container.Database.Client); Assert.Equal(containerName, container.Id); Assert.Equal(expectedEndpoint, container.Database.Client.Endpoint.ToString()); + Assert.False(container.Database.Client.ClientOptions.LimitToEndpoint); } [Fact] - public void AddAzureCosmosDatabase_CreatesNewClient_WhenConfigureClientOptionsProvided() + public void AddAzureCosmosDatabase_NoAddKeyedContainer_AddsDatabase() { var builder = Host.CreateEmptyApplicationBuilder(null); var databaseName = "testdb"; @@ -362,108 +388,248 @@ public void AddAzureCosmosDatabase_CreatesNewClient_WhenConfigureClientOptionsPr PopulateConfiguration(builder.Configuration, connectionString); - builder.AddAzureCosmosClient("cosmos"); - builder.AddAzureCosmosDatabase("cosmos", configureClientOptions: options => - { - options.LimitToEndpoint = false; - }); + builder.AddAzureCosmosDatabase("cosmos"); using var host = builder.Build(); + var database = host.Services.GetService(); + var client = host.Services.GetService(); + var container = host.Services.GetService(); - var client = host.Services.GetRequiredService(); - var database = host.Services.GetRequiredService(); - - Assert.NotSame(client, database.Client); + Assert.NotNull(database); Assert.Equal(databaseName, database.Id); - Assert.Equal(expectedEndpoint, database.Client.Endpoint.ToString()); - Assert.False(database.Client.ClientOptions.LimitToEndpoint); + Assert.Null(client); + Assert.Null(container); } [Fact] - public void AddAzureCosmosContainer_CreatesNewClient_WhenConfigureClientOptionsProvided() + public void AddAzureCosmosDatabase_AddKeyedContainer_RegistersContainerWithDatabaseKey() { var builder = Host.CreateEmptyApplicationBuilder(null); var databaseName = "testdb"; var containerName = "testcontainer"; var expectedEndpoint = "https://localhost:8081/"; - var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};Container={containerName};"; + var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};"; - PopulateConfiguration(builder.Configuration, connectionString); + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString), + new KeyValuePair("ConnectionStrings:container1", $"{connectionString}Container={containerName};") + ]); - builder.AddAzureCosmosClient("cosmos"); - builder.AddAzureCosmosContainer("cosmos", configureClientOptions: options => - { - options.LimitToEndpoint = false; - }); + var databaseBuilder = builder.AddAzureCosmosDatabase("cosmos"); + databaseBuilder.AddKeyedContainer("container1"); using var host = builder.Build(); - var client = host.Services.GetRequiredService(); - var container = host.Services.GetRequiredService(); + // Database and client should not be registered + var database = host.Services.GetService(); + var client = host.Services.GetService(); + Assert.Null(client); - Assert.NotSame(client, container.Database.Client); + // Verify that database was registered + Assert.NotNull(database); + Assert.Equal(databaseName, database.Id); + + // Verify container was registered with the key correct key + var container = host.Services.GetRequiredKeyedService("container1"); + Assert.NotNull(container); Assert.Equal(containerName, container.Id); - Assert.Equal(expectedEndpoint, container.Database.Client.Endpoint.ToString()); - Assert.False(container.Database.Client.ClientOptions.LimitToEndpoint); } [Fact] - public void AddKeyedAzureCosmosDatabase_CreatesNewClient_WhenConfigureClientOptionsProvided() + public void AddAzureCosmosDatabase_AddMultipleContainers_RegistersAllWithSameClient() { var builder = Host.CreateEmptyApplicationBuilder(null); - var serviceKey = "cosmos-key"; var databaseName = "testdb"; + var container1Name = "container1"; + var container2Name = "container2"; var expectedEndpoint = "https://localhost:8081/"; var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};"; - PopulateConfiguration(builder.Configuration, connectionString, serviceKey); + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString), + new KeyValuePair("ConnectionStrings:container1", $"{connectionString}Container={container1Name};"), + new KeyValuePair("ConnectionStrings:container2", $"{connectionString}Container={container2Name};") + ]); - builder.AddKeyedAzureCosmosClient(serviceKey); - builder.AddKeyedAzureCosmosDatabase(serviceKey, configureClientOptions: options => - { - options.LimitToEndpoint = false; - }); + builder.AddAzureCosmosDatabase("cosmos") + .AddKeyedContainer("container1") + .AddKeyedContainer("container2"); using var host = builder.Build(); - var client = host.Services.GetRequiredKeyedService(serviceKey); - var database = host.Services.GetRequiredKeyedService(serviceKey); + var container1 = host.Services.GetRequiredKeyedService("container1"); + var container2 = host.Services.GetRequiredKeyedService("container2"); - Assert.NotSame(client, database.Client); - Assert.Equal(databaseName, database.Id); - Assert.Equal(expectedEndpoint, database.Client.Endpoint.ToString()); - Assert.False(database.Client.ClientOptions.LimitToEndpoint); + // Different containers + Assert.NotNull(container1); + Assert.NotNull(container2); + Assert.Equal(container1Name, container1.Id); + Assert.Equal(container2Name, container2.Id); + + // With the same client + Assert.Same(container2.Database.Client, container1.Database.Client); } [Fact] - public void AddKeyedAzureCosmosContainer_CreatesNewClient_WhenConfigureClientOptionsProvided() + public void AddAzureCosmosDatabase_AddKeyedContainer_ThrowsWhenContainerNameMissing() + { + var builder = Host.CreateEmptyApplicationBuilder(null); + var databaseName = "testdb"; + var connectionString = $"AccountEndpoint=https://localhost:8081/;AccountKey=fake;Database={databaseName};"; + + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString), + new KeyValuePair("ConnectionStrings:container1", connectionString) + ]); + + builder.AddAzureCosmosDatabase("cosmos") + .AddKeyedContainer("container1"); + + using var host = builder.Build(); + + var exception = Assert.Throws( + () => host.Services.GetRequiredKeyedService("container1")); + Assert.Contains("A Container could not be configured", exception.Message); + } + + [Fact] + public void AddAzureCosmosDatabase_AddKeyedContainer_WorksWithNoConnectionString() + { + var builder = Host.CreateEmptyApplicationBuilder(null); + var databaseName = "testdb"; + var connectionString = $"AccountEndpoint=https://localhost:8081/;AccountKey=fake;Database={databaseName};"; + + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString) + ]); + + builder.AddAzureCosmosDatabase("cosmos") + .AddKeyedContainer("container1"); + + using var host = builder.Build(); + + var container = host.Services.GetKeyedService("container1"); + var database = host.Services.GetRequiredService(); + + Assert.NotNull(container); + Assert.Equal("container1", container.Id); + Assert.Equal(databaseName, container.Database.Id); + Assert.Equal("https://localhost:8081/", container.Database.Client.Endpoint.ToString()); + } + + [Fact] + public void AddAzureCosmosDatabase_AddKeyedContainer_CustomizeClientOptions() { var builder = Host.CreateEmptyApplicationBuilder(null); - var serviceKey = "cosmos-key"; var databaseName = "testdb"; var containerName = "testcontainer"; var expectedEndpoint = "https://localhost:8081/"; - var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};Container={containerName};"; + var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={databaseName};"; - PopulateConfiguration(builder.Configuration, connectionString, serviceKey); + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString), + new KeyValuePair("ConnectionStrings:container1", $"{connectionString}Container={containerName};") + ]); - builder.AddKeyedAzureCosmosClient(serviceKey); - builder.AddKeyedAzureCosmosContainer(serviceKey, configureClientOptions: options => - { - options.LimitToEndpoint = false; - }); + builder.AddAzureCosmosDatabase("cosmos", + configureClientOptions: options => { + options.ApplicationName = "TestApp"; + options.LimitToEndpoint = false; + }) + .AddKeyedContainer("container1"); using var host = builder.Build(); - var client = host.Services.GetRequiredKeyedService(serviceKey); - var container = host.Services.GetRequiredKeyedService(serviceKey); - - Assert.NotSame(client, container.Database.Client); - Assert.Equal(containerName, container.Id); - Assert.Equal(expectedEndpoint, container.Database.Client.Endpoint.ToString()); + // Verify container has the expected client options + var container = host.Services.GetRequiredKeyedService("container1"); + Assert.Contains("TestApp", container.Database.Client.ClientOptions.ApplicationName); Assert.False(container.Database.Client.ClientOptions.LimitToEndpoint); } + [Fact] + public void AddAzureCosmosDatabase_ConfigureSettings_AppliesToAllContainers() + { + var builder = Host.CreateEmptyApplicationBuilder(null); + var databaseName = "testdb"; + var container1Name = "container1"; + var container2Name = "container2"; + var expectedEndpoint = "https://localhost:8081/"; + var connectionString = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;"; + + builder.Configuration.AddInMemoryCollection([ + new KeyValuePair("ConnectionStrings:cosmos", connectionString), + new KeyValuePair("ConnectionStrings:container1", $"{connectionString}Container={container1Name};"), + new KeyValuePair("ConnectionStrings:container2", $"{connectionString}Container={container2Name};") + ]); + + var databaseBuilder = builder.AddAzureCosmosDatabase("cosmos", + configureSettings: settings => + { + // Database name comes from settings, not connection string + settings.DatabaseName = databaseName; + settings.DisableTracing = true; + }) + .AddKeyedContainer("container1") + .AddKeyedContainer("container2"); + + using var host = builder.Build(); + + var container1 = host.Services.GetRequiredKeyedService("container1"); + var container2 = host.Services.GetRequiredKeyedService("container2"); + + Assert.Equal(databaseName, container1.Database.Id); + Assert.Equal(databaseName, container2.Database.Id); + Assert.Same(container1.Database.Client, container2.Database.Client); + } + + [Fact] + public void AddAzureCosmosDatabase_CalledMultipleTimes_CreatesIndependentBuilders() + { + var builder = Host.CreateEmptyApplicationBuilder(null); + var database1Name = "db1"; + var database2Name = "db2"; + var container1Name = "users"; + var container2Name = "orders"; + var container3Name = "products"; + var expectedEndpoint = "https://localhost:8081/"; + var connectionString1 = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={database1Name};"; + var connectionString2 = $"AccountEndpoint={expectedEndpoint};AccountKey=fake;Database={database2Name};"; + + builder.Configuration.AddInMemoryCollection([ + // First database connection + new KeyValuePair("ConnectionStrings:cosmos1", connectionString1), + new KeyValuePair("ConnectionStrings:users", $"{connectionString1}Container={container1Name};"), + new KeyValuePair("ConnectionStrings:orders", $"{connectionString1}Container={container2Name};"), + + // Second database connection + new KeyValuePair("ConnectionStrings:cosmos2", connectionString2), + new KeyValuePair("ConnectionStrings:products", $"{connectionString2}Container={container3Name};") + ]); + + // Create two separate database builders + builder.AddAzureCosmosDatabase("cosmos1") + .AddKeyedContainer("users") + .AddKeyedContainer("orders"); + builder.AddAzureCosmosDatabase("cosmos2") + .AddKeyedContainer("products"); + + using var host = builder.Build(); + + var usersContainer = host.Services.GetRequiredKeyedService(container1Name); + var ordersContainer = host.Services.GetRequiredKeyedService(container2Name); + var productsContainer = host.Services.GetRequiredKeyedService(container3Name); + + Assert.Equal(container1Name, usersContainer.Id); + Assert.Equal(container2Name, ordersContainer.Id); + Assert.Equal(container3Name, productsContainer.Id); + Assert.Equal(database1Name, usersContainer.Database.Id); + Assert.Equal(database1Name, ordersContainer.Database.Id); + Assert.Equal(database2Name, productsContainer.Database.Id); + + Assert.Same(usersContainer.Database.Client, ordersContainer.Database.Client); + Assert.NotSame(usersContainer.Database.Client, productsContainer.Database.Client); + } + private static void PopulateConfiguration(ConfigurationManager configuration, string connectionString, string? key = null) => configuration.AddInMemoryCollection([ new KeyValuePair($"ConnectionStrings:{key ?? "cosmos"}", connectionString)