Add RemoveExtension and RemoveDbContext APIs for removing provider configuration (#37891)

Fixes #35126

Co-authored-by: AndriySvyryd <6539701+AndriySvyryd@users.noreply.github.com>

Author: Copilot

src/EFCore/DbContextOptions.cs

@@ -93,6 +93,15 @@ public virtual TExtension GetExtension<TExtension>()
     public abstract DbContextOptions WithExtension<TExtension>(TExtension extension)
         where TExtension : class, IDbContextOptionsExtension;
 
+    /// <summary>
+    ///     Removes the given extension from the underlying options and creates a new
+    ///     <see cref="DbContextOptions" /> with the extension removed.
+    /// </summary>
+    /// <typeparam name="TExtension">The type of extension to be removed.</typeparam>
+    /// <returns>The new options instance with the extension removed.</returns>
+    public abstract DbContextOptions WithoutExtension<TExtension>()
+        where TExtension : class, IDbContextOptionsExtension;
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore/DbContextOptionsBuilder.cs

@@ -786,6 +786,17 @@ public virtual DbContextOptionsBuilder UseAsyncSeeding(Func<DbContext, bool, Can
     void IDbContextOptionsBuilderInfrastructure.AddOrUpdateExtension<TExtension>(TExtension extension)
         => _options = _options.WithExtension(extension);
 
+    /// <summary>
+    ///     Removes the extension of the given type from the options. If no extension of the given type exists, this is a no-op.
+    /// </summary>
+    /// <remarks>
+    ///     This method is intended for use by extension methods to configure the context. It is not intended to be used in
+    ///     application code.
+    /// </remarks>
+    /// <typeparam name="TExtension">The type of extension to be removed.</typeparam>
+    void IDbContextOptionsBuilderInfrastructure.RemoveExtension<TExtension>()
+        => _options = _options.WithoutExtension<TExtension>();
+
     private DbContextOptionsBuilder WithOption(Func<CoreOptionsExtension, CoreOptionsExtension> withFunc)
     {
         ((IDbContextOptionsBuilderInfrastructure)this).AddOrUpdateExtension(

src/EFCore/DbContextOptions`.cs

@@ -57,6 +57,30 @@ public override DbContextOptions WithExtension<TExtension>(TExtension extension)
         return new DbContextOptions<TContext>(ExtensionsMap.SetItem(type, (extension, ordinal)));
     }
 
+    /// <inheritdoc />
+    public override DbContextOptions WithoutExtension<TExtension>()
+    {
+        var type = typeof(TExtension);
+        if (!ExtensionsMap.TryGetValue(type, out var removedValue))
+        {
+            return this;
+        }
+
+        var removedOrdinal = removedValue.Ordinal;
+        var newMap = ExtensionsMap.Remove(type);
+
+        // Renormalize ordinals for extensions that followed the removed one
+        foreach (var (key, value) in newMap)
+        {
+            if (value.Ordinal > removedOrdinal)
+            {
+                newMap = newMap.SetItem(key, (value.Extension, value.Ordinal - 1));
+            }
+        }
+
+        return new DbContextOptions<TContext>(newMap);
+    }
+
     /// <summary>
     ///     The type of context that these options are for (<typeparamref name="TContext" />).
     /// </summary>

src/EFCore/Extensions/EntityFrameworkServiceCollectionExtensions.cs

@@ -1163,6 +1163,66 @@ public static IServiceCollection ConfigureDbContext
         return serviceCollection;
     }
 
+    /// <summary>
+    ///     Removes services for the given context type from the <see cref="IServiceCollection" />.
+    /// </summary>
+    /// <remarks>
+    ///     <para>
+    ///         This method can be used to remove the context registration in integration testing scenarios
+    ///         where a different database provider is used for tests.
+    ///     </para>
+    ///     <para>
+    ///         See <see href="https://aka.ms/efcore-docs-di">Using DbContext with dependency injection</see> for more information and examples.
+    ///     </para>
+    /// </remarks>
+    /// <typeparam name="TContext">The type of context to be removed.</typeparam>
+    /// <param name="serviceCollection">The <see cref="IServiceCollection" /> to remove services from.</param>
+    /// <param name="removeConfigurationOnly">
+    ///     If <see langword="true" />, only the <see cref="IDbContextOptionsConfiguration{TContext}" /> registrations will be removed;
+    ///     the context itself will remain registered. If <see langword="false" /> (the default), all services related to the context
+    ///     will be removed.
+    /// </param>
+    /// <returns>The same service collection so that multiple calls can be chained.</returns>
+    public static IServiceCollection RemoveDbContext
+        <[DynamicallyAccessedMembers(DbContext.DynamicallyAccessedMemberTypes)] TContext>(
+            this IServiceCollection serviceCollection,
+            bool removeConfigurationOnly = false)
+        where TContext : DbContext
+    {
+        Check.NotNull(serviceCollection);
+
+        if (removeConfigurationOnly)
+        {
+            var configurations = serviceCollection
+                .Where(d => d.ServiceType == typeof(IDbContextOptionsConfiguration<TContext>))
+                .ToList();
+
+            foreach (var descriptor in configurations)
+            {
+                serviceCollection.Remove(descriptor);
+            }
+        }
+        else
+        {
+            var descriptorsToRemove = serviceCollection
+                .Where(d => d.ServiceType == typeof(TContext)
+                    || d.ServiceType == typeof(DbContextOptions<TContext>)
+                    || d.ServiceType == typeof(IDbContextOptionsConfiguration<TContext>)
+                    || d.ServiceType == typeof(IDbContextFactorySource<TContext>)
+                    || d.ServiceType == typeof(IDbContextFactory<TContext>)
+                    || d.ServiceType == typeof(IDbContextPool<TContext>)
+                    || d.ServiceType == typeof(IScopedDbContextLease<TContext>))
+                .ToList();
+
+            foreach (var descriptor in descriptorsToRemove)
+            {
+                serviceCollection.Remove(descriptor);
+            }
+        }
+
+        return serviceCollection;
+    }
+
     private static void AddCoreServices<TContextImplementation>(
         IServiceCollection serviceCollection,
         Action<IServiceProvider, DbContextOptionsBuilder>? optionsAction,

src/EFCore/Infrastructure/IDbContextOptionsBuilderInfrastructure.cs

@@ -36,4 +36,21 @@ public interface IDbContextOptionsBuilderInfrastructure
     /// <param name="extension">The extension to be added.</param>
     void AddOrUpdateExtension<TExtension>(TExtension extension)
         where TExtension : class, IDbContextOptionsExtension;
+
+    /// <summary>
+    ///     <para>
+    ///         Removes the extension of the given type from the options. If no extension of the given type exists, this is a no-op.
+    ///     </para>
+    ///     <para>
+    ///         This method is intended for use by extension methods to configure the context. It is not intended to be used in
+    ///         application code.
+    ///     </para>
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-providers">Implementation of database providers and extensions</see>
+    ///     for more information and examples.
+    /// </remarks>
+    /// <typeparam name="TExtension">The type of extension to be removed.</typeparam>
+    void RemoveExtension<TExtension>()
+        where TExtension : class, IDbContextOptionsExtension;
 }

test/EFCore.Tests/DbContextOptionsTest.cs

@@ -96,6 +96,76 @@ public void Can_update_an_existing_extension()
         Assert.Same(extension2, optionsBuilder.Options.FindExtension<FakeDbContextOptionsExtension1>());
     }
 
+    [ConditionalFact]
+    public void Can_remove_an_existing_extension()
+    {
+        var optionsBuilder = new DbContextOptionsBuilder();
+
+        var extension1 = new FakeDbContextOptionsExtension1();
+        var extension2 = new FakeDbContextOptionsExtension2();
+
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension1);
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension2);
+
+        Assert.Equal(2, optionsBuilder.Options.Extensions.Count());
+
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).RemoveExtension<FakeDbContextOptionsExtension1>();
+
+        Assert.Single(optionsBuilder.Options.Extensions);
+        Assert.Null(optionsBuilder.Options.FindExtension<FakeDbContextOptionsExtension1>());
+        Assert.Same(extension2, optionsBuilder.Options.FindExtension<FakeDbContextOptionsExtension2>());
+    }
+
+    [ConditionalFact]
+    public void Removing_non_existent_extension_is_no_op()
+    {
+        var optionsBuilder = new DbContextOptionsBuilder();
+
+        var extension = new FakeDbContextOptionsExtension1();
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension);
+
+        Assert.Single(optionsBuilder.Options.Extensions);
+
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).RemoveExtension<FakeDbContextOptionsExtension2>();
+
+        Assert.Single(optionsBuilder.Options.Extensions);
+        Assert.Same(extension, optionsBuilder.Options.FindExtension<FakeDbContextOptionsExtension1>());
+    }
+
+    [ConditionalFact]
+    public void Removing_extension_from_middle_renormalizes_ordinals_and_preserves_insertion_order()
+    {
+        var optionsBuilder = new DbContextOptionsBuilder();
+
+        var extension1 = new FakeDbContextOptionsExtension1();
+        var extension2 = new FakeDbContextOptionsExtension2();
+        var extension3 = new FakeDbContextOptionsExtension3();
+
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension1);
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension2);
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension3);
+
+        Assert.Equal(3, optionsBuilder.Options.Extensions.Count());
+
+        // Remove the middle extension
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).RemoveExtension<FakeDbContextOptionsExtension2>();
+
+        Assert.Equal(2, optionsBuilder.Options.Extensions.Count());
+        var extensionsList = optionsBuilder.Options.Extensions.ToList();
+        Assert.Same(extension1, extensionsList[0]);
+        Assert.Same(extension3, extensionsList[1]);
+
+        // Add a new extension after removing the middle one - ordinals should stay contiguous
+        var extension2New = new FakeDbContextOptionsExtension2();
+        ((IDbContextOptionsBuilderInfrastructure)optionsBuilder).AddOrUpdateExtension(extension2New);
+
+        Assert.Equal(3, optionsBuilder.Options.Extensions.Count());
+        extensionsList = optionsBuilder.Options.Extensions.ToList();
+        Assert.Same(extension1, extensionsList[0]);
+        Assert.Same(extension3, extensionsList[1]);
+        Assert.Same(extension2New, extensionsList[2]);
+    }
+
     [ConditionalFact]
     public void IsConfigured_returns_true_if_any_provider_extensions_have_been_added()
     {
@@ -199,6 +269,42 @@ public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
         }
     }
 
+    private class FakeDbContextOptionsExtension3 : IDbContextOptionsExtension
+    {
+        private DbContextOptionsExtensionInfo _info;
+
+        public DbContextOptionsExtensionInfo Info
+            => _info ??= new ExtensionInfo(this);
+
+        public bool AppliedServices { get; private set; }
+
+        public virtual void ApplyServices(IServiceCollection services)
+            => AppliedServices = true;
+
+        public virtual void Validate(IDbContextOptions options)
+        {
+        }
+
+        private sealed class ExtensionInfo(IDbContextOptionsExtension extension) : DbContextOptionsExtensionInfo(extension)
+        {
+            public override bool IsDatabaseProvider
+                => false;
+
+            public override int GetServiceProviderHashCode()
+                => 0;
+
+            public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo other)
+                => true;
+
+            public override string LogFragment
+                => "";
+
+            public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
+            {
+            }
+        }
+    }
+
     [ConditionalFact]
     public void UseModel_on_generic_builder_returns_generic_builder()
     {