CosmosDB with Managed Identity

CustomCosmosExtension.cs

using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;

namespace CosmosTest.Extensions;

public class CustomCosmosExtension : CosmosOptionsExtension
{
    public string ManagedIdentityClientId { get; private set; } = string.Empty;
    public override void ApplyServices(IServiceCollection services)
    {
        // This is the first reason this class is created.
        // To be able to add the singleton with credentials.
        services.AddEntityFrameworkCosmosCustom();
    }

    public CustomCosmosExtension(CustomCosmosExtension ext) : base(ext)
    {
        ManagedIdentityClientId = ext.ManagedIdentityClientId;
    }

    public CustomCosmosExtension()
    {
        
    }

    protected override CosmosOptionsExtension Clone()
    {
        return new CustomCosmosExtension(this);
    }

    // This is the second reason this class is created.
    // To store client id for passthrough to other classes
    public void WithManagedIdentityClientId(string clientId)
    {
        ManagedIdentityClientId = clientId;
    }
}

CustomCosmosServiceCollectionExtension.cs

using System.ComponentModel;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal;
using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Conventions.Internal;
using Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
using Microsoft.EntityFrameworkCore.Cosmos.ValueGeneration.Internal;
using Microsoft.EntityFrameworkCore.Diagnostics;
using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Metadata.Conventions.Infrastructure;
using Microsoft.EntityFrameworkCore.Query;
using Microsoft.EntityFrameworkCore.Storage;
using Microsoft.EntityFrameworkCore.ValueGeneration;

namespace CosmosTest.Extensions;

public static class CustomCosmosServiceCollectionExtensions
{
    public static IServiceCollection AddCosmosCustom<TContext>(
        this IServiceCollection serviceCollection,
        string endpoint,
        string databaseName,
        string managedIdentityClientId,
        Action<CosmosDbContextOptionsBuilder>? cosmosOptionsAction = null,
        Action<DbContextOptionsBuilder>? optionsAction = null)
        where TContext : DbContext
    {
        return serviceCollection.AddDbContext<TContext>(
            (serviceProvider, options) =>
            {
                optionsAction?.Invoke(options);
                options.UseCosmosCustom(endpoint, databaseName, managedIdentityClientId, cosmosOptionsAction);
            });
    }

    [EditorBrowsable(EditorBrowsableState.Never)]
    public static IServiceCollection AddEntityFrameworkCosmosCustom(this IServiceCollection serviceCollection)
    {
        // Here in the custom implementation the singleton gets added. This is needed because this is where the credentials are gathered.
        // It also does TryAdd in the default implementation, so this custom singleton does not get overridden
        // Marked with //Custom

        var builder = new EntityFrameworkServicesBuilder(serviceCollection)
            .TryAdd<LoggingDefinitions, CosmosLoggingDefinitions>()
            .TryAdd<IDatabaseProvider, DatabaseProvider<CosmosOptionsExtension>>()
            .TryAdd<IDatabase, CosmosDatabaseWrapper>()
            .TryAdd<IExecutionStrategyFactory, CosmosExecutionStrategyFactory>()
            .TryAdd<IDbContextTransactionManager, CosmosTransactionManager>()
            .TryAdd<IModelValidator, CosmosModelValidator>()
            .TryAdd<IProviderConventionSetBuilder, CosmosConventionSetBuilder>()
            .TryAdd<IValueGeneratorSelector, CosmosValueGeneratorSelector>()
            .TryAdd<IDatabaseCreator, CosmosDatabaseCreator>()
            .TryAdd<IQueryContextFactory, CosmosQueryContextFactory>()
            .TryAdd<ITypeMappingSource, CosmosTypeMappingSource>()
            .TryAdd<IQueryableMethodTranslatingExpressionVisitorFactory, CosmosQueryableMethodTranslatingExpressionVisitorFactory>()
            .TryAdd<IShapedQueryCompilingExpressionVisitorFactory, CosmosShapedQueryCompilingExpressionVisitorFactory>()
            .TryAdd<IQueryTranslationPreprocessorFactory, CosmosQueryTranslationPreprocessorFactory>()
            .TryAdd<IQueryCompilationContextFactory, CosmosQueryCompilationContextFactory>()
            .TryAdd<IQueryTranslationPostprocessorFactory, CosmosQueryTranslationPostprocessorFactory>()
            //Custom
            .TryAdd<ISingletonOptions, ICustomCosmosSingletonOptions>(p => p.GetRequiredService<ICustomCosmosSingletonOptions>())
            .TryAddProviderSpecificServices(
                b => b
                    //Custom
                    .TryAddSingleton<ICustomCosmosSingletonOptions, CustomCosmosSingletonOptions>()
                    //Custom
                    .TryAddSingleton<ISingletonCosmosClientWrapper, CustomSingletonCosmosClientWrapper>()
                    .TryAddSingleton<IQuerySqlGeneratorFactory, QuerySqlGeneratorFactory>()
                    .TryAddScoped<ISqlExpressionFactory, SqlExpressionFactory>()
                    .TryAddScoped<IMemberTranslatorProvider, CosmosMemberTranslatorProvider>()
                    .TryAddScoped<IMethodCallTranslatorProvider, CosmosMethodCallTranslatorProvider>()
                    .TryAddScoped<ICosmosClientWrapper, CosmosClientWrapper>()
            );

        builder.TryAddCoreServices();

        return serviceCollection;
    }
}

CustomCosmosSingletonOptions.cs

using System.Diagnostics.CodeAnalysis;
using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
using Microsoft.EntityFrameworkCore.Infrastructure;

namespace CosmosTest.Extensions;

public class CustomCosmosSingletonOptions : CosmosSingletonOptions, ICustomCosmosSingletonOptions
{
    public string ManagedClientIdentityId { get; private set; } = string.Empty;

    public override void Initialize(IDbContextOptions options)
    {
        var cosmosOptions = options.FindExtension<CustomCosmosExtension>() ?? throw new ArgumentNullException(nameof(CustomCosmosExtension));
        ManagedClientIdentityId = cosmosOptions.ManagedIdentityClientId;
        base.Initialize(options);
    }
}

CustomSingletonCosmosClientWrapper.cs

using Azure.Identity;
using Microsoft.Azure.Cosmos;
using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
using Microsoft.EntityFrameworkCore.Infrastructure;

namespace CosmosTest.Extensions;

public class CustomSingletonCosmosClientWrapper : ISingletonCosmosClientWrapper
{
    private CosmosClient? _client;
    private readonly string? _endpoint;
    private static readonly string UserAgent = " Microsoft.EntityFrameworkCore.Cosmos/" + ProductInfo.GetVersion();
    private readonly CosmosClientOptions _options;
    private readonly string _managedIdentityClientId;

    public CustomSingletonCosmosClientWrapper(ICustomCosmosSingletonOptions options)
    {
        _endpoint = options.AccountEndpoint;
        _managedIdentityClientId = options.ManagedClientIdentityId;
        var configuration = new CosmosClientOptions { ApplicationName = UserAgent, Serializer = new JsonCosmosSerializer() };
        
        // Because _options is private we cannot reuse them from parent
        // and need to copy
        if (options.Region != null)
        {
            configuration.ApplicationRegion = options.Region;
        }

        if (options.LimitToEndpoint != null)
        {
            configuration.LimitToEndpoint = options.LimitToEndpoint.Value;
        }

        if (options.ConnectionMode != null)
        {
            configuration.ConnectionMode = options.ConnectionMode.Value;
        }

        if (options.WebProxy != null)
        {
            configuration.WebProxy = options.WebProxy;
        }

        if (options.RequestTimeout != null)
        {
            configuration.RequestTimeout = options.RequestTimeout.Value;
        }

        if (options.OpenTcpConnectionTimeout != null)
        {
            configuration.OpenTcpConnectionTimeout = options.OpenTcpConnectionTimeout.Value;
        }

        if (options.IdleTcpConnectionTimeout != null)
        {
            configuration.IdleTcpConnectionTimeout = options.IdleTcpConnectionTimeout.Value;
        }

        if (options.GatewayModeMaxConnectionLimit != null)
        {
            configuration.GatewayModeMaxConnectionLimit = options.GatewayModeMaxConnectionLimit.Value;
        }

        if (options.MaxTcpConnectionsPerEndpoint != null)
        {
            configuration.MaxTcpConnectionsPerEndpoint = options.MaxTcpConnectionsPerEndpoint.Value;
        }

        if (options.MaxRequestsPerTcpConnection != null)
        {
            configuration.MaxRequestsPerTcpConnection = options.MaxRequestsPerTcpConnection.Value;
        }

        if (options.HttpClientFactory != null)
        {
            configuration.HttpClientFactory = options.HttpClientFactory;
        }

        _options = configuration;
    }

    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    protected virtual void Dispose(bool disposing)
    {
        _client?.Dispose();
        _client = null;
    }

    public CosmosClient Client
        => _client ?? CreateClient();


    private CosmosClient CreateClient()
    {
        var credentials = new DefaultAzureCredential(new DefaultAzureCredentialOptions()
        {
            ManagedIdentityClientId = _managedIdentityClientId,
        });

        var client = new CosmosClient(_endpoint, credentials, _options);
        _client = client;
        return client;
    }
}

ICustomCosmosSingletonOptions.cs

using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;

namespace CosmosTest.Extensions;

// Extend interface with extra client id
public interface ICustomCosmosSingletonOptions : ICosmosSingletonOptions
{
    string ManagedClientIdentityId { get; }
}

OptionsExtension.cs

using CosmosTest.Extensions;
using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
using Microsoft.EntityFrameworkCore.Infrastructure;

// ReSharper disable once CheckNamespace
namespace Microsoft.EntityFrameworkCore;

public static class OptionsExtension
{
    public static DbContextOptionsBuilder UseCosmosCustom(
        this DbContextOptionsBuilder optionsBuilder,
        string accountEndpoint,
        string databaseName,
        string clientId,
        Action<CosmosDbContextOptionsBuilder>? cosmosOptionsAction = null)
    {
        var customExtension = optionsBuilder.Options.FindExtension<CustomCosmosExtension>() as CosmosOptionsExtension
                                                ?? new CustomCosmosExtension();

        ((CustomCosmosExtension)customExtension)
            .WithManagedIdentityClientId(clientId);
        customExtension = customExtension
            .WithAccountEndpoint(accountEndpoint)
            .WithDatabaseName(databaseName);

        var defaultExtension = optionsBuilder.Options.FindExtension<CosmosOptionsExtension>()
                                          ?? new CosmosOptionsExtension();

        defaultExtension = defaultExtension
            .WithAccountEndpoint(accountEndpoint)
            .WithDatabaseName(databaseName);


        // There are hardcoded dependencies in EF Core Cosmos internally to CosmosOptionsExtension. So you HAVE to register it.
        // The generic parameters are not used under the hood unfortunately (in DbContextOptions<TContext> it does GetType()).
        // If they were used, we could create a cleaner solution.
        // Now we register both, first one for the custom credentials, and second one for the hardcoded dependencies
        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(customExtension);
        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(defaultExtension);

        cosmosOptionsAction?.Invoke(new CosmosDbContextOptionsBuilder(optionsBuilder));

        return optionsBuilder;
    }
}

Program.cs

...
var builder = WebApplication.CreateBuilder(args);

builder.Services.AddCosmosCustom<CosmosContext>(
    endpoint: "https://westus22.documents.azure.com:443/",
    databaseName: "MyDb",
    managedIdentityClientId: "90df9398-990d-459f-8833-bfa4d762a4d7");

...

If you're on the hunt for top-notch managed IT services in New Jersey, I highly recommend checking out RedPaladin. Their comprehensive suite of managed IT solutions https://redpaladin.com/services/managed-it-services-new-jersey/ is tailored to meet the diverse needs of businesses in the region. From proactive network monitoring to cybersecurity measures and beyond, RedPaladin excels in providing reliable and efficient IT support. Give them a look if you're in need of dependable IT services for your business.