From ef89493924f025c7d643be76fdbcdc9fd3fec007 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Sun, 2 Aug 2026 08:53:25 -0700 Subject: [PATCH 1/2] Add configuration support for Cosmos grain storage Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: e82f060d-d09b-42f9-8b0f-34da38b386fa --- .../CosmosGrainStorageProviderBuilder.cs | 52 +++++++++++++++++++ .../CosmosHostingExtensionsTests.cs | 46 ++++++++++++++++ 2 files changed, 98 insertions(+) create mode 100644 src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs diff --git a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs new file mode 100644 index 00000000000..c415a3ac28b --- /dev/null +++ b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs @@ -0,0 +1,52 @@ +using Microsoft.Extensions.Configuration; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using Orleans.Persistence.Cosmos; +using Orleans.Providers; + +[assembly: RegisterProvider("AzureCosmosDB", "GrainStorage", "Silo", typeof(CosmosGrainStorageProviderBuilder))] + +namespace Orleans.Hosting; + +internal sealed class CosmosGrainStorageProviderBuilder : IProviderBuilder +{ + public void Configure(ISiloBuilder builder, string? name, IConfigurationSection configurationSection) + { + ArgumentException.ThrowIfNullOrEmpty(name); + + var documentIdProviderKey = configurationSection["DocumentIdProviderKey"]; + if (!string.IsNullOrEmpty(documentIdProviderKey)) + { + builder.Services.AddKeyedSingleton( + name, + (services, _) => services.GetRequiredKeyedService(documentIdProviderKey)); + } + + builder.AddCosmosGrainStorage(name, (OptionsBuilder optionsBuilder) => + { + optionsBuilder.Bind(configurationSection); + optionsBuilder.Configure((options, services) => + { + var serviceKey = configurationSection["ServiceKey"]; + if (!string.IsNullOrEmpty(serviceKey)) + { + options.ConfigureCosmosClient( + provider => new ValueTask(provider.GetRequiredKeyedService(serviceKey))); + return; + } + + var connectionName = configurationSection["ConnectionName"]; + var connectionString = configurationSection["ConnectionString"]; + if (!string.IsNullOrEmpty(connectionName) && string.IsNullOrEmpty(connectionString)) + { + connectionString = services.GetRequiredService().GetConnectionString(connectionName); + } + + if (!string.IsNullOrEmpty(connectionString)) + { + options.ConfigureCosmosClient(connectionString); + } + }); + }); + } +} diff --git a/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs b/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs index f4d7cebb824..ccb5ac0eedc 100644 --- a/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs +++ b/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs @@ -1,5 +1,7 @@ +using Microsoft.Extensions.Configuration; using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Options; using Orleans.Configuration; using Orleans.Hosting; using Orleans.Persistence.Cosmos; @@ -24,6 +26,39 @@ public void AddCosmosGrainStorage_ProviderTypesAreRegisteredByKey() Assert.IsType(host.Services.GetRequiredKeyedService("second")); } + [Fact] + public void CosmosGrainStorageProviderBuilder_UsesConfiguredDocumentIdProviderFromDI() + { + const string storageName = "configured-storage"; + const string documentIdProviderKey = "custom-document-ids"; + var configuration = new ConfigurationBuilder() + .AddInMemoryCollection(new Dictionary + { + ["Cosmos:DatabaseName"] = "configured-database", + ["Cosmos:DocumentIdProviderKey"] = documentIdProviderKey + }) + .Build(); + + using var host = new HostBuilder() + .UseOrleans(builder => + { + builder.Services.AddSingleton(); + builder.Services.AddKeyedSingleton(documentIdProviderKey); + new CosmosGrainStorageProviderBuilder().Configure(builder, storageName, configuration.GetSection("Cosmos")); + }) + .Build(); + + var configuredProvider = host.Services.GetRequiredKeyedService(storageName); + var registeredProvider = host.Services.GetRequiredKeyedService(documentIdProviderKey); + var options = host.Services.GetRequiredService>().Get(storageName); + + Assert.Same(registeredProvider, configuredProvider); + Assert.Same( + host.Services.GetRequiredService(), + Assert.IsType(configuredProvider).Dependency); + Assert.Equal("configured-database", options.DatabaseName); + } + #pragma warning disable CS0618 // Type or member is obsolete [Fact] public void AddCosmosGrainStorage_LegacyPartitionKeyProvidersAreRegisteredByKey() @@ -76,6 +111,17 @@ private sealed class SecondDocumentIdProvider : IDocumentIdProvider public ValueTask<(string DocumentId, string PartitionKey)> GetDocumentIdentifiers(string grainType, GrainId grainId) => default; } + private sealed class ConfiguredDocumentIdProvider(DocumentIdProviderDependency dependency) : IDocumentIdProvider + { + public DocumentIdProviderDependency Dependency { get; } = dependency; + + public ValueTask<(string DocumentId, string PartitionKey)> GetDocumentIdentifiers(string grainType, GrainId grainId) => default; + } + + private sealed class DocumentIdProviderDependency + { + } + #pragma warning disable CS0618 // Type or member is obsolete private sealed class FirstPartitionKeyProvider : IPartitionKeyProvider { From 5d85d0f2330c58e36e549e37ce21e4cc8b9923b0 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Tue, 4 Aug 2026 14:15:29 -0700 Subject: [PATCH 2/2] fix(cosmos): report missing named connection strings Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: cc82c70d-1c36-4b98-92d4-d9e0c2bce556 --- .../CosmosGrainStorageProviderBuilder.cs | 4 ++++ .../CosmosHostingExtensionsTests.cs | 23 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs index c415a3ac28b..078f2b48d88 100644 --- a/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs +++ b/src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs @@ -40,6 +40,10 @@ public void Configure(ISiloBuilder builder, string? name, IConfigurationSection if (!string.IsNullOrEmpty(connectionName) && string.IsNullOrEmpty(connectionString)) { connectionString = services.GetRequiredService().GetConnectionString(connectionName); + if (string.IsNullOrEmpty(connectionString)) + { + throw new InvalidOperationException($"Connection string '{connectionName}' was not found."); + } } if (!string.IsNullOrEmpty(connectionString)) diff --git a/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs b/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs index ccb5ac0eedc..b4c2473b78c 100644 --- a/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs +++ b/test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs @@ -59,6 +59,29 @@ public void CosmosGrainStorageProviderBuilder_UsesConfiguredDocumentIdProviderFr Assert.Equal("configured-database", options.DatabaseName); } + [Fact] + public void CosmosGrainStorageProviderBuilder_ThrowsWhenNamedConnectionStringIsMissing() + { + const string storageName = "configured-storage"; + const string connectionName = "missing"; + var configuration = new ConfigurationBuilder() + .AddInMemoryCollection(new Dictionary + { + ["Cosmos:ConnectionName"] = connectionName + }) + .Build(); + + using var host = new HostBuilder() + .UseOrleans(builder => + new CosmosGrainStorageProviderBuilder().Configure(builder, storageName, configuration.GetSection("Cosmos"))) + .Build(); + + var options = host.Services.GetRequiredService>(); + var exception = Assert.Throws(() => options.Get(storageName)); + + Assert.Equal($"Connection string '{connectionName}' was not found.", exception.Message); + } + #pragma warning disable CS0618 // Type or member is obsolete [Fact] public void AddCosmosGrainStorage_LegacyPartitionKeyProvidersAreRegisteredByKey()