From ff3874cffb1a9ac8949f5dc1fa7336e58b35a055 Mon Sep 17 00:00:00 2001
From: OmnipotentOwl <1769881+OmnipotentOwl@users.noreply.github.com>
Date: Sat, 26 Oct 2024 19:12:40 -0400
Subject: [PATCH] Add CosmosClusteringProviderBuilder for Orleans clustering

Introduce CosmosClusteringProviderBuilder class implementing IProviderBuilder<ISiloBuilder> and IProviderBuilder<IClientBuilder> interfaces. Register the class as a provider for both "Silo" and "Client" roles. Implement Configure methods to set up Cosmos DB clustering options and client configuration. Import necessary namespaces to support the new functionality.

This adds support for using the Cosmos DB Clustering Provider from Orleans IConfiguration setup.
---
 .../CosmosClusteringProviderBuilder.cs        | 115 ++++++++++++++++++
 1 file changed, 115 insertions(+)
 create mode 100644 src/Azure/Orleans.Clustering.Cosmos/CosmosClusteringProviderBuilder.cs

diff --git a/src/Azure/Orleans.Clustering.Cosmos/CosmosClusteringProviderBuilder.cs b/src/Azure/Orleans.Clustering.Cosmos/CosmosClusteringProviderBuilder.cs
new file mode 100644
index 0000000000..87d343474f
--- /dev/null
+++ b/src/Azure/Orleans.Clustering.Cosmos/CosmosClusteringProviderBuilder.cs
@@ -0,0 +1,115 @@
+using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
+using Orleans;
+using Orleans.Clustering.Cosmos;
+using Orleans.Hosting;
+using Orleans.Providers;
+
+[assembly: RegisterProvider("Cosmos", "Clustering", "Silo", typeof(CosmosClusteringProviderBuilder))]
+[assembly: RegisterProvider("Cosmos", "Clustering", "Client", typeof(CosmosClusteringProviderBuilder))]
+
+namespace Orleans.Hosting;
+
+internal sealed class CosmosClusteringProviderBuilder : IProviderBuilder<ISiloBuilder>, IProviderBuilder<IClientBuilder>
+{
+    public void Configure(ISiloBuilder builder, string name, IConfigurationSection configurationSection) =>
+        builder.UseCosmosClustering(optionsBuilder =>
+            optionsBuilder.Configure<IServiceProvider>((options, services) =>
+            {
+                var databaseName = configurationSection[nameof(options.DatabaseName)];
+                if (!string.IsNullOrEmpty(databaseName))
+                {
+                    options.DatabaseName = databaseName;
+                }
+                var containerName = configurationSection[nameof(options.ContainerName)];
+                if (!string.IsNullOrEmpty(containerName))
+                {
+                    options.ContainerName = containerName;
+                }
+                if (bool.TryParse(configurationSection[nameof(options.IsResourceCreationEnabled)], out var irce))
+                {
+                    options.IsResourceCreationEnabled = irce;
+                }
+                if(int.TryParse(configurationSection[nameof(options.DatabaseThroughput)], out var dt))
+                {
+                    options.DatabaseThroughput = dt;
+                }
+                if(bool.TryParse(configurationSection[nameof(options.CleanResourcesOnInitialization)], out var croi))
+                {
+                    options.CleanResourcesOnInitialization = croi;
+                }
+
+                var serviceKey = configurationSection["ServiceKey"];
+                if(!string.IsNullOrEmpty(serviceKey))
+                {
+                    options.ConfigureCosmosClient(sp=>
+                        new ValueTask<CosmosClient>(sp.GetRequiredKeyedService<CosmosClient>(serviceKey)));
+                }
+                else
+                {
+                    var connectionName = configurationSection["ConnectionName"];
+                    var connectionString = configurationSection["ConnectionString"];
+                    if (!string.IsNullOrEmpty(connectionName) && string.IsNullOrEmpty(connectionString))
+                    {
+                        var rootConfiguration = services.GetRequiredService<IConfiguration>();
+                        connectionString = rootConfiguration.GetConnectionString(connectionName);
+                    }
+
+                    if (!string.IsNullOrEmpty(connectionString))
+                    {
+                        options.ConfigureCosmosClient(connectionString);
+                    }
+                }
+            }));
+
+    public void Configure(IClientBuilder builder, string name, IConfigurationSection configurationSection) =>
+        builder.UseCosmosGatewayListProvider(optionsBuilder =>
+            optionsBuilder.Configure<IServiceProvider>((options, services) =>
+            {
+                var databaseName = configurationSection[nameof(options.DatabaseName)];
+                if (!string.IsNullOrEmpty(databaseName))
+                {
+                    options.DatabaseName = databaseName;
+                }
+                var containerName = configurationSection[nameof(options.ContainerName)];
+                if (!string.IsNullOrEmpty(containerName))
+                {
+                    options.ContainerName = containerName;
+                }
+                if (bool.TryParse(configurationSection[nameof(options.IsResourceCreationEnabled)], out var irce))
+                {
+                    options.IsResourceCreationEnabled = irce;
+                }
+                if (int.TryParse(configurationSection[nameof(options.DatabaseThroughput)], out var dt))
+                {
+                    options.DatabaseThroughput = dt;
+                }
+                if (bool.TryParse(configurationSection[nameof(options.CleanResourcesOnInitialization)], out var croi))
+                {
+                    options.CleanResourcesOnInitialization = croi;
+                }
+
+                var serviceKey = configurationSection["ServiceKey"];
+                if (!string.IsNullOrEmpty(serviceKey))
+                {
+                    options.ConfigureCosmosClient(sp =>
+                        new ValueTask<CosmosClient>(sp.GetRequiredKeyedService<CosmosClient>(serviceKey)));
+                }
+                else
+                {
+                    // Construct a connection multiplexer from a connection string.
+                    var connectionName = configurationSection["ConnectionName"];
+                    var connectionString = configurationSection["ConnectionString"];
+                    if (!string.IsNullOrEmpty(connectionName) && string.IsNullOrEmpty(connectionString))
+                    {
+                        var rootConfiguration = services.GetRequiredService<IConfiguration>();
+                        connectionString = rootConfiguration.GetConnectionString(connectionName);
+                    }
+
+                    if (!string.IsNullOrEmpty(connectionString))
+                    {
+                        options.ConfigureCosmosClient(connectionString);
+                    }
+                }
+            }));
+}
\ No newline at end of file