Introduce enumerator-like API to allow peeking ahead to know whether …

…more batches are available

src/EFCore.Relational/Infrastructure/EntityFrameworkRelationalServicesBuilder.cs

@@ -60,7 +60,8 @@ public static readonly IDictionary<Type, ServiceCharacteristics> RelationalServi
             { typeof(IRawSqlCommandBuilder), new ServiceCharacteristics(ServiceLifetime.Singleton) },
             { typeof(IQuerySqlGeneratorFactory), new ServiceCharacteristics(ServiceLifetime.Singleton) },
             { typeof(IModificationCommandFactory), new ServiceCharacteristics(ServiceLifetime.Singleton) },
-            { typeof(ICommandBatchPreparer), new ServiceCharacteristics(ServiceLifetime.Scoped) },
+            // TODO: Can be singleton?
+            { typeof(ICommandBatchPreparerFactory), new ServiceCharacteristics(ServiceLifetime.Scoped) },
             { typeof(IModificationCommandBatchFactory), new ServiceCharacteristics(ServiceLifetime.Scoped) },
             { typeof(IRelationalSqlTranslatingExpressionVisitorFactory), new ServiceCharacteristics(ServiceLifetime.Scoped) },
             { typeof(IMethodCallTranslatorProvider), new ServiceCharacteristics(ServiceLifetime.Scoped) },
@@ -139,7 +140,7 @@ public override EntityFrameworkServicesBuilder TryAddCoreServices()
         TryAdd<IValueGeneratorSelector, RelationalValueGeneratorSelector>();
         TryAdd<IRelationalCommandBuilderFactory, RelationalCommandBuilderFactory>();
         TryAdd<IRawSqlCommandBuilder, RawSqlCommandBuilder>();
-        TryAdd<ICommandBatchPreparer, CommandBatchPreparer>();
+        TryAdd<ICommandBatchPreparerFactory, CommandBatchPreparerFactory>();
         TryAdd<IModificationCommandFactory, ModificationCommandFactory>();
         TryAdd<IMigrationsModelDiffer, MigrationsModelDiffer>();
         TryAdd<IMigrationsSqlGenerator, MigrationsSqlGenerator>();

src/EFCore.Relational/Migrations/Internal/MigrationsModelDiffer.cs

@@ -2186,10 +2186,10 @@ private IEnumerable<MigrationOperation> GetDataOperations(
         }
 
         var model = updateAdapter.Model.GetRelationalModel();
-        var commandBatches = new CommandBatchPreparer(CommandBatchPreparerDependencies)
-            .BatchCommands(entries, updateAdapter);
+        var commandBatchPreparerFactory = new CommandBatchPreparerFactory(CommandBatchPreparerDependencies);
+        var commandBatchPreparer = commandBatchPreparerFactory.CreateCommandBatchPreparer(entries, updateAdapter);
 
-        foreach (var commandBatch in commandBatches)
+        while (commandBatchPreparer.TryGetNextBatch(out var commandBatch))
         {
             InsertDataOperation? batchInsertOperation = null;
             foreach (var command in commandBatch.ModificationCommands)

src/EFCore.Relational/Storage/RelationalDatabase.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.Update.Internal;
+
 namespace Microsoft.EntityFrameworkCore.Storage;
 
 /// <summary>
@@ -51,7 +53,7 @@ public RelationalDatabase(
     /// <returns>The number of state entries persisted to the database.</returns>
     public override int SaveChanges(IList<IUpdateEntry> entries)
         => RelationalDependencies.BatchExecutor.Execute(
-            RelationalDependencies.BatchPreparer.BatchCommands(
+            RelationalDependencies.CommandBatchPreparerFactory.CreateCommandBatchPreparer(
                 entries,
                 Dependencies.UpdateAdapterFactory.Create()),
             RelationalDependencies.Connection);
@@ -70,7 +72,7 @@ public override Task<int> SaveChangesAsync(
         IList<IUpdateEntry> entries,
         CancellationToken cancellationToken = default)
         => RelationalDependencies.BatchExecutor.ExecuteAsync(
-            RelationalDependencies.BatchPreparer.BatchCommands(
+            RelationalDependencies.CommandBatchPreparerFactory.CreateCommandBatchPreparer(
                 entries,
                 Dependencies.UpdateAdapterFactory.Create()),
             RelationalDependencies.Connection,

src/EFCore.Relational/Storage/RelationalDatabaseDependencies.cs

@@ -46,19 +46,19 @@ public sealed record RelationalDatabaseDependencies
     /// </remarks>
     [EntityFrameworkInternal]
     public RelationalDatabaseDependencies(
-        ICommandBatchPreparer batchPreparer,
+        ICommandBatchPreparerFactory commandBatchPreparerFactory,
         IBatchExecutor batchExecutor,
         IRelationalConnection connection)
     {
-        BatchPreparer = batchPreparer;
+        CommandBatchPreparerFactory = commandBatchPreparerFactory;
         BatchExecutor = batchExecutor;
         Connection = connection;
     }
 
     /// <summary>
     ///     The <see cref="ICommandBatchPreparer" /> to be used.
     /// </summary>
-    public ICommandBatchPreparer BatchPreparer { get; init; }
+    public ICommandBatchPreparerFactory CommandBatchPreparerFactory { get; init; }
 
     /// <summary>
     ///     The <see cref="IBatchExecutor" /> to be used.

src/EFCore.Relational/Update/IBatchExecutor.cs

@@ -28,17 +28,17 @@ public interface IBatchExecutor
     /// <summary>
     ///     Executes the commands in the batches against the given database connection.
     /// </summary>
-    /// <param name="commandBatches">The batches to execute.</param>
+    /// <param name="commandBatchPreparer">A service to provide the the batches to execute.</param>
     /// <param name="connection">The database connection to use.</param>
     /// <returns>The total number of rows affected.</returns>
     int Execute(
-        IEnumerable<ModificationCommandBatch> commandBatches,
+        ICommandBatchPreparer commandBatchPreparer,
         IRelationalConnection connection);
 
     /// <summary>
     ///     Executes the commands in the batches against the given database connection.
     /// </summary>
-    /// <param name="commandBatches">The batches to execute.</param>
+    /// <param name="commandBatchPreparer">A service to provide the the batches to execute.</param>
     /// <param name="connection">The database connection to use.</param>
     /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
     /// <returns>
@@ -47,7 +47,7 @@ int Execute(
     /// </returns>
     /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
     Task<int> ExecuteAsync(
-        IEnumerable<ModificationCommandBatch> commandBatches,
+        ICommandBatchPreparer commandBatchPreparer,
         IRelationalConnection connection,
         CancellationToken cancellationToken = default);
 }

src/EFCore.Relational/Update/ICommandBatchPreparer.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics.CodeAnalysis;
+
 namespace Microsoft.EntityFrameworkCore.Update;
 
 /// <summary>
@@ -27,13 +29,18 @@ namespace Microsoft.EntityFrameworkCore.Update;
 public interface ICommandBatchPreparer
 {
     /// <summary>
-    ///     Creates the command batches needed to insert/update/delete the entities represented by the given
-    ///     list of <see cref="IUpdateEntry" />s.
+    ///     Attempts to returns the next batch to be executed.
+    /// </summary>
+    /// <param name="batch">The next batch.</param>
+    /// <returns>
+    ///     <see langword="true" /> if the next batch was populated into <paramref name="batch" />, or <see langword="false" /> if there
+    ///     were no more batched.
+    /// </returns>
+    bool TryGetNextBatch([NotNullWhen(true)] out ModificationCommandBatch? batch);
+
+    /// <summary>
+    ///     Returns whether there are more batched to be executed.
     /// </summary>
-    /// <param name="entries">The entries that represent the entities to be modified.</param>
-    /// <param name="updateAdapter">The model data.</param>
-    /// <returns>The list of batches to execute.</returns>
-    IEnumerable<ModificationCommandBatch> BatchCommands(
-        IList<IUpdateEntry> entries,
-        IUpdateAdapter updateAdapter);
+    /// <returns>Whether there are more batched to be executed.</returns>
+    bool HasNextBatch();
 }

src/EFCore.Relational/Update/ICommandBatchPreparerFactory.cs

@@ -0,0 +1,33 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Update;
+
+/// <summary>
+///     <para>
+///         A service for creating <see cref="ICommandBatchPreparer" /> instances.
+///     </para>
+///     <para>
+///         This type is typically used by database providers; it is generally not used in application code.
+///     </para>
+/// </summary>
+/// <remarks>
+///     <para>
+///         The service lifetime is <see cref="ServiceLifetime.Singleton" />. This means a single instance
+///         is used by many <see cref="DbContext" /> instances. The implementation must be thread-safe.
+///         This service cannot depend on services registered as <see cref="ServiceLifetime.Scoped" />.
+///     </para>
+///     <para>
+///         See <see href="https://aka.ms/efcore-docs-providers">Implementation of database providers and extensions</see>
+///         for more information and examples.
+///     </para>
+/// </remarks>
+public interface ICommandBatchPreparerFactory
+{
+    /// <summary>
+    ///     Creates a new instance of <see cref="ICommandBatchPreparer" />.
+    /// </summary>
+    /// <returns>A new <see cref="ICommandBatchPreparer" /> instance.</returns>
+    // TODO: Doc
+    ICommandBatchPreparer CreateCommandBatchPreparer(IList<IUpdateEntry> entries, IUpdateAdapter updateAdapter);
+}

src/EFCore.Relational/Update/Internal/BatchExecutor.cs

@@ -49,19 +49,14 @@ public BatchExecutor(
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual int Execute(
-        IEnumerable<ModificationCommandBatch> commandBatches,
+        ICommandBatchPreparer commandBatchPreparer,
         IRelationalConnection connection)
     {
-        using var batchEnumerator = commandBatches.GetEnumerator();
-
-        if (!batchEnumerator.MoveNext())
+        if (!commandBatchPreparer.TryGetNextBatch(out var batch))
         {
             return 0;
         }
 
-        var currentBatch = batchEnumerator.Current;
-        var nextBatch = batchEnumerator.MoveNext() ? batchEnumerator.Current : null;
-
         var rowsAffected = 0;
         var transaction = connection.CurrentTransaction;
         var beganTransaction = false;
@@ -74,7 +69,7 @@ public virtual int Execute(
                 && transactionEnlistManager?.CurrentAmbientTransaction is null
                 && CurrentContext.Context.Database.AutoTransactionsEnabled
                 // Don't start a transaction if we have a single batch which doesn't require a transaction (single command), for perf.
-                && (nextBatch is not null || currentBatch.RequiresTransaction))
+                && (commandBatchPreparer.HasNextBatch() || batch.RequiresTransaction))
             {
                 transaction = connection.BeginTransaction();
                 beganTransaction = true;
@@ -91,14 +86,12 @@ public virtual int Execute(
                 }
             }
 
-            while (currentBatch is not null)
+            do
             {
-                currentBatch.Execute(connection);
-                rowsAffected += currentBatch.ModificationCommands.Count;
-
-                currentBatch = nextBatch;
-                nextBatch = batchEnumerator.MoveNext() ? batchEnumerator.Current : null;
+                batch.Execute(connection);
+                rowsAffected += batch.ModificationCommands.Count;
             }
+            while (commandBatchPreparer.TryGetNextBatch(out batch));
 
             if (beganTransaction)
             {
@@ -158,20 +151,15 @@ public virtual int Execute(
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual async Task<int> ExecuteAsync(
-        IEnumerable<ModificationCommandBatch> commandBatches,
+        ICommandBatchPreparer commandBatchPreparer,
         IRelationalConnection connection,
         CancellationToken cancellationToken = default)
     {
-        using var batchEnumerator = commandBatches.GetEnumerator();
-
-        if (!batchEnumerator.MoveNext())
+        if (!commandBatchPreparer.TryGetNextBatch(out var batch))
         {
             return 0;
         }
 
-        var currentBatch = batchEnumerator.Current;
-        var nextBatch = batchEnumerator.MoveNext() ? batchEnumerator.Current : null;
-
         var rowsAffected = 0;
         var transaction = connection.CurrentTransaction;
         var beganTransaction = false;
@@ -184,7 +172,7 @@ public virtual async Task<int> ExecuteAsync(
                 && transactionEnlistManager?.CurrentAmbientTransaction is null
                 && CurrentContext.Context.Database.AutoTransactionsEnabled
                 // Don't start a transaction if we have a single batch which doesn't require a transaction (single command), for perf.
-                && (nextBatch is not null || currentBatch.RequiresTransaction))
+                && (commandBatchPreparer.HasNextBatch() || batch.RequiresTransaction))
             {
                 transaction = await connection.BeginTransactionAsync(cancellationToken).ConfigureAwait(false);
                 beganTransaction = true;
@@ -201,14 +189,12 @@ public virtual async Task<int> ExecuteAsync(
                 }
             }
 
-            while (currentBatch is not null)
+            do
             {
-                await currentBatch.ExecuteAsync(connection, cancellationToken).ConfigureAwait(false);
-                rowsAffected += currentBatch.ModificationCommands.Count;
-
-                currentBatch = nextBatch;
-                nextBatch = batchEnumerator.MoveNext() ? batchEnumerator.Current : null;
+                await batch.ExecuteAsync(connection, cancellationToken).ConfigureAwait(false);
+                rowsAffected += batch.ModificationCommands.Count;
             }
+            while (commandBatchPreparer.TryGetNextBatch(out batch));
 
             if (beganTransaction)
             {