Pomelo.EFCoreでSharding

resolver.cs

public class ConnectionStringResolver : IConnectionStringResolver<MyDbContext, long>
{
    public string GetConnectionStringForKey(long key)
    {
        return "connection-string";
    }
}

sharding.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Threading.Tasks;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Storage.Internal;
using Microsoft.Extensions.DependencyInjection;
using MySqlConnector;
using Pomelo.EntityFrameworkCore.MySql.Infrastructure.Internal;

namespace Sharding
{
    public interface IConnectionStringResolver<TContext, TKey>
    {
        string GetConnectionStringForKey(TKey key);
    }

    public interface IShardedDbContextResolver<out TContext, in TKey> : IDisposable, IAsyncDisposable
    {
        TContext GetShard(TKey key);
    }

    public class ShardedDbContextResolver<TContextService, TContextImplementation, TKey> : IShardedDbContextResolver<TContextService, TKey>
        where TContextImplementation : DbContext, TContextService
        where TKey : notnull
    {
        private readonly IServiceProvider _sp;

        private readonly object _gate = new object();

        private readonly Dictionary<TKey, TContextImplementation> _shards = new Dictionary<TKey, TContextImplementation>();

        private readonly IConnectionStringResolver<TContextImplementation, TKey> _connectionStringResolver;
        private readonly Action<IServiceProvider, DbContextOptionsBuilder>? _optionsAction;

        private bool _disposed = false;

        public ShardedDbContextResolver(IServiceProvider sp,
            IConnectionStringResolver<TContextImplementation, TKey> connectionStringResolver,
            Action<IServiceProvider, DbContextOptionsBuilder>? optionsAction)
        {
            _sp = sp;
            _connectionStringResolver = connectionStringResolver;
            _optionsAction = optionsAction;
        }

        public TContextService GetShard(TKey key)
        {
            lock(_gate)
            {
                TContextImplementation? result;
                if (_shards.TryGetValue(key, out result))
                {
                    return result;
                }
                else
                {
                    result = Create(key, _connectionStringResolver, _optionsAction);
                    _shards.Add(key, result);
                    return result;
                }
            }
        }
#pragma warning disable EF1001
        private TContextImplementation Create(TKey key,
            IConnectionStringResolver<TContextImplementation, TKey> connectionStringResolver, 
            Action<IServiceProvider, DbContextOptionsBuilder>? optionsAction)
        {
            var builder = new DbContextOptionsBuilder<TContextImplementation>(
                new DbContextOptions<TContextImplementation>(
                    new Dictionary<Type, IDbContextOptionsExtension>()));

            builder.UseApplicationServiceProvider(_sp);

            optionsAction?.Invoke(_sp, builder); //ServerVersionだけ設定される想定

            var connectionString = connectionStringResolver.GetConnectionStringForKey(key);

            var resolvedConnectionString = new NamedConnectionStringResolver(builder.Options)
                .ResolveConnectionString(connectionString);

            var csb = new MySqlConnectionStringBuilder(resolvedConnectionString)
            {
                AllowUserVariables = true,
                UseAffectedRows = false
            };

            resolvedConnectionString = csb.ConnectionString;

            var extension = (MySqlOptionsExtension)GetOrCreateExtension(builder)
                .WithConnectionString(resolvedConnectionString);

            ((IDbContextOptionsBuilderInfrastructure)builder).AddOrUpdateExtension(extension);

            return ActivatorUtilities.CreateInstance<TContextImplementation>(_sp, builder.Options);
        }

        private static MySqlOptionsExtension GetOrCreateExtension(DbContextOptionsBuilder optionsBuilder)
            => optionsBuilder.Options.FindExtension<MySqlOptionsExtension>()
                ?? new MySqlOptionsExtension();
#pragma warning restore EF1001

        public void Dispose()
        {
            if (_disposed)
                return;
            
            _disposed = true;
            lock(_gate)
            {
                foreach(var pair in _shards)
                {
                    pair.Value.Dispose();
                }
                _shards.Clear();
            }
        }

        public ValueTask DisposeAsync()
        {
            //TODO:
            return ValueTask.CompletedTask;
        }
    }
}

namespace Microsoft.Extensions.DependencyInjection
{
    using Sharding;
    using Microsoft.Extensions.DependencyInjection.Extensions;

    public static class ShardingServiceCollectionExtensions
    {
        public static IServiceCollection AddShardedDbContext<TContext, TKey, TConnectionStringResolver>(this IServiceCollection serviceCollection,
            Action<IServiceProvider, DbContextOptionsBuilder>? optionsAction,
            ServiceLifetime contextLifetime = ServiceLifetime.Scoped)
            where TContext : DbContext
            where TConnectionStringResolver : class, IConnectionStringResolver<TContext, TKey>
            where TKey : notnull
            => serviceCollection.AddShardedDbContext<TContext, TContext, TKey, TConnectionStringResolver>(optionsAction, contextLifetime);

        public static IServiceCollection AddShardedDbContext<TContextService, TContextImplementation, TKey, TConnectionStringResolver>(
            this IServiceCollection serviceCollection,
            Action<IServiceProvider, DbContextOptionsBuilder>? optionsAction,
            ServiceLifetime contextLifetime = ServiceLifetime.Scoped)
            where TContextImplementation : DbContext, TContextService
            where TConnectionStringResolver : class, IConnectionStringResolver<TContextImplementation, TKey>
            where TKey : notnull
        {
            if (optionsAction != null)
            {
                CheckContextConstructors<TContextImplementation>();
            }

            serviceCollection.TryAdd(new ServiceDescriptor(
                typeof(IConnectionStringResolver<TContextImplementation, TKey>),
                typeof(TConnectionStringResolver),
                contextLifetime) //singletonにしたり別引数でもらってもいいかもしれない
            );

            serviceCollection.TryAdd(new ServiceDescriptor(
                typeof(IShardedDbContextResolver<TContextService, TKey>),
                sp => ActivatorUtilities.CreateInstance(sp, typeof(ShardedDbContextResolver<TContextService, TContextImplementation, TKey>), optionsAction!),
                contextLifetime)
            );

            return serviceCollection;
        }

        private static void CheckContextConstructors<TContext>()
            where TContext : DbContext
        {
            var declaredConstructors = typeof(TContext).GetTypeInfo().DeclaredConstructors.ToList();
            if (declaredConstructors.Count == 1
                && declaredConstructors[0].GetParameters().Length == 0)
            {
                //throw new ArgumentException(CoreStrings.DbContextMissingConstructor(typeof(TContext).ShortDisplayName()));
                throw new ArgumentException(typeof(TContext).Name); //ちょっと雑だが……
            }
        }
    }
}

usage.cs

services.AddShardedDbContext<MyDbContext, long, ConnectionStringResolver>((sp, b) => b.UseMySql(
    new MySqlServerVersion(new Version(8, 0))
    )
);