Skip to content

feat(cosmos): add configuration support for grain storage #10326

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ReubenBond merged 2 commits into from
Aug 4, 2026
Merged

feat(cosmos): add configuration support for grain storage #10326

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ef89493
Add configuration support for Cosmos grain storage
ReubenBond Aug 2, 2026
5d85d0f
fix(cosmos): report missing named connection strings
ReubenBond Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions src/Azure/Orleans.Persistence.Cosmos/CosmosGrainStorageProviderBuilder.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
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<ISiloBuilder>
{
public void Configure(ISiloBuilder builder, string? name, IConfigurationSection configurationSection)
{
ArgumentException.ThrowIfNullOrEmpty(name);

var documentIdProviderKey = configurationSection["DocumentIdProviderKey"];
if (!string.IsNullOrEmpty(documentIdProviderKey))
{
builder.Services.AddKeyedSingleton<IDocumentIdProvider>(
name,
(services, _) => services.GetRequiredKeyedService<IDocumentIdProvider>(documentIdProviderKey));
}

builder.AddCosmosGrainStorage(name, (OptionsBuilder<CosmosGrainStorageOptions> optionsBuilder) =>
{
optionsBuilder.Bind(configurationSection);
optionsBuilder.Configure<IServiceProvider>((options, services) =>
{
var serviceKey = configurationSection["ServiceKey"];
if (!string.IsNullOrEmpty(serviceKey))
{
options.ConfigureCosmosClient(
provider => new ValueTask<CosmosClient>(provider.GetRequiredKeyedService<CosmosClient>(serviceKey)));
return;
}

var connectionName = configurationSection["ConnectionName"];
var connectionString = configurationSection["ConnectionString"];
if (!string.IsNullOrEmpty(connectionName) && string.IsNullOrEmpty(connectionString))
{
connectionString = services.GetRequiredService<IConfiguration>().GetConnectionString(connectionName);
if (string.IsNullOrEmpty(connectionString))
{
throw new InvalidOperationException($"Connection string '{connectionName}' was not found.");
}
}

if (!string.IsNullOrEmpty(connectionString))
{
options.ConfigureCosmosClient(connectionString);
}
});
});
}
}
69 changes: 69 additions & 0 deletions test/Extensions/Orleans.Cosmos.Tests/CosmosHostingExtensionsTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -24,6 +26,62 @@ public void AddCosmosGrainStorage_ProviderTypesAreRegisteredByKey()
Assert.IsType<SecondDocumentIdProvider>(host.Services.GetRequiredKeyedService<IDocumentIdProvider>("second"));
}

[Fact]
public void CosmosGrainStorageProviderBuilder_UsesConfiguredDocumentIdProviderFromDI()
{
const string storageName = "configured-storage";
const string documentIdProviderKey = "custom-document-ids";
var configuration = new ConfigurationBuilder()
.AddInMemoryCollection(new Dictionary<string, string?>
{
["Cosmos:DatabaseName"] = "configured-database",
["Cosmos:DocumentIdProviderKey"] = documentIdProviderKey
})
.Build();

using var host = new HostBuilder()
.UseOrleans(builder =>
{
builder.Services.AddSingleton<DocumentIdProviderDependency>();
builder.Services.AddKeyedSingleton<IDocumentIdProvider, ConfiguredDocumentIdProvider>(documentIdProviderKey);
new CosmosGrainStorageProviderBuilder().Configure(builder, storageName, configuration.GetSection("Cosmos"));
})
.Build();

var configuredProvider = host.Services.GetRequiredKeyedService<IDocumentIdProvider>(storageName);
var registeredProvider = host.Services.GetRequiredKeyedService<IDocumentIdProvider>(documentIdProviderKey);
var options = host.Services.GetRequiredService<IOptionsMonitor<CosmosGrainStorageOptions>>().Get(storageName);

Assert.Same(registeredProvider, configuredProvider);
Assert.Same(
host.Services.GetRequiredService<DocumentIdProviderDependency>(),
Assert.IsType<ConfiguredDocumentIdProvider>(configuredProvider).Dependency);
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<string, string?>
{
["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<IOptionsMonitor<CosmosGrainStorageOptions>>();
var exception = Assert.Throws<InvalidOperationException>(() => 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()
Expand Down Expand Up @@ -76,6 +134,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
{
Expand Down
Loading