AddData and RemoveData for DbContext

src/Directory.Build.props

@@ -2,7 +2,7 @@
 <Project>
   <PropertyGroup>
     <NoWarn>CS1591;CA1416;CS8632</NoWarn>
-    <Version>15.3.0</Version>
+    <Version>15.4.0</Version>
     <AssemblyVersion>1.0.0</AssemblyVersion>
     <ContinuousIntegrationBuild>false</ContinuousIntegrationBuild>
     <CheckEolTargetFramework>false</CheckEolTargetFramework>

src/EfLocalDb/DbContextExtensions.cs

@@ -0,0 +1,64 @@
+namespace EfLocalDb;
+
+public static partial class DbContextExtensions
+{
+    internal static IEnumerable<object> ExpandEnumerable(IEnumerable<object> entities, IReadOnlyList<IEntityType> entityTypes)
+    {
+        foreach (var entity in entities)
+        {
+            if (entity is IEnumerable enumerable)
+            {
+                var entityType = entity.GetType();
+                if (entityTypes.Any(_ => _.ClrType != entityType))
+                {
+                    foreach (var nested in enumerable)
+                    {
+                        yield return nested;
+                    }
+
+                    continue;
+                }
+            }
+
+            yield return entity;
+        }
+    }
+
+    public static Task AddData<TDbContext>(this TDbContext context, IEnumerable<object> entities)
+        where TDbContext : DbContext =>
+        context.AddData(entities, context.Model.GetEntityTypes().ToArray());
+
+    public static Task AddData<TDbContext>(this TDbContext context, params object[] entities)
+        where TDbContext : DbContext =>
+        context.AddData((IEnumerable<object>) entities);
+
+    internal static Task AddData<TDbContext>(this TDbContext context, IEnumerable<object> entities, IReadOnlyList<IEntityType> entityTypes)
+        where TDbContext : DbContext
+    {
+        foreach (var entity in ExpandEnumerable(entities, entityTypes))
+        {
+            context.Add(entity);
+        }
+
+        return context.SaveChangesAsync();
+    }
+
+    public static Task RemoveData<TDbContext>(this TDbContext context, IEnumerable<object> entities)
+        where TDbContext : DbContext =>
+        context.RemoveData(entities, context.Model.GetEntityTypes().ToArray());
+
+    public static Task RemoveData<TDbContext>(this TDbContext context, params object[] entities)
+        where TDbContext : DbContext =>
+        context.RemoveData((IEnumerable<object>) entities);
+
+    internal static Task RemoveData<TDbContext>(this TDbContext context, IEnumerable<object> entities, IReadOnlyList<IEntityType> entityTypes)
+        where TDbContext : DbContext
+    {
+        foreach (var entity in ExpandEnumerable(entities, entityTypes))
+        {
+            context.Remove(entity);
+        }
+
+        return context.SaveChangesAsync();
+    }
+}

src/EfLocalDb/SqlDatabase.cs

@@ -141,25 +141,6 @@ public async Task Delete()
     public DbSet<T> Set<T>()
         where T : class => NoTrackingContext.Set<T>();
 
-    IEnumerable<object> ExpandEnumerable(IEnumerable<object> entities)
-    {
-        foreach (var entity in entities)
-        {
-            if (entity is IEnumerable enumerable)
-            {
-                var entityType = entity.GetType();
-                if (instance.EntityTypes.Any(_ => _.ClrType != entityType))
-                {
-                    foreach (var nested in enumerable)
-                    {
-                        yield return nested;
-                    }
-
-                    continue;
-                }
-            }
-
-            yield return entity;
-        }
-    }
+    IEnumerable<object> ExpandEnumerable(IEnumerable<object> entities) =>
+        DbContextExtensions.ExpandEnumerable(entities, instance.EntityTypes);
 }

src/EfLocalDb/SqlDatabase_Add.cs

@@ -2,24 +2,16 @@
 
 public partial class SqlDatabase<TDbContext>
 {
-    public Task AddData(IEnumerable<object> entities) => Add(entities, Context);
+    public Task AddData(IEnumerable<object> entities) =>
+        Context.AddData(entities, instance.EntityTypes);
 
-    Task Add(IEnumerable<object> entities, TDbContext context)
-    {
-        foreach (var entity in ExpandEnumerable(entities))
-        {
-            context.Add(entity);
-        }
-
-        return context.SaveChangesAsync();
-    }
-
-    public Task AddData(params object[] entities) => AddData((IEnumerable<object>) entities);
+    public Task AddData(params object[] entities) =>
+        AddData((IEnumerable<object>) entities);
 
     public async Task AddDataUntracked(IEnumerable<object> entities)
     {
         await using var context = NewDbContext();
-        await Add(entities, context);
+        await context.AddData(entities, instance.EntityTypes);
     }
 
     public Task AddDataUntracked(params object[] entities) => AddDataUntracked((IEnumerable<object>) entities);

src/EfLocalDb/SqlDatabase_Remove.cs

@@ -2,24 +2,16 @@
 
 public partial class SqlDatabase<TDbContext>
 {
-    public Task RemoveData(IEnumerable<object> entities) => Remove(entities, Context);
+    public Task RemoveData(IEnumerable<object> entities) =>
+        Context.RemoveData(entities, instance.EntityTypes);
 
-    Task Remove(IEnumerable<object> entities, TDbContext context)
-    {
-        foreach (var entity in ExpandEnumerable(entities))
-        {
-            context.Remove(entity);
-        }
-
-        return context.SaveChangesAsync();
-    }
-
-    public Task RemoveData(params object[] entities) => RemoveData((IEnumerable<object>) entities);
+    public Task RemoveData(params object[] entities) =>
+        RemoveData((IEnumerable<object>) entities);
 
     public async Task RemoveDataUntracked(IEnumerable<object> entities)
     {
         await using var context = NewDbContext();
-        await Remove(entities, context);
+        await context.RemoveData(entities, instance.EntityTypes);
     }
 
     public Task RemoveDataUntracked(params object[] entities) => RemoveDataUntracked((IEnumerable<object>) entities);