Add test for parameter managed across batches (#28721)

Closes #23281
dotnet · Aug 15, 2022 · adf8eb4 · adf8eb4
1 parent 99c1aa2
commit adf8eb4
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 17 deletions.
diff --git a/src/EFCore.Relational/Update/ReaderModificationCommandBatch.cs b/src/EFCore.Relational/Update/ReaderModificationCommandBatch.cs
@@ -121,21 +121,18 @@ public override bool TryAddCommand(IReadOnlyModificationCommand modificationComm
             return true;
         }
 
-        RollbackLastCommand();
+        Check.DebugAssert(ReferenceEquals(modificationCommand, _modificationCommands[^1]),
+            "ReferenceEquals(modificationCommand, _modificationCommands[^1])");
 
-        // The command's column modifications had their parameter names generated, that needs to be rolled back as well.
-        foreach (var columnModification in modificationCommand.ColumnModifications)
-        {
-            columnModification.ResetParameterNames();
-        }
+        RollbackLastCommand(modificationCommand);
 
         return false;
     }
 
     /// <summary>
     ///     Rolls back the last command added. Used when adding a command caused the batch to become invalid (e.g. CommandText too long).
     /// </summary>
-    protected virtual void RollbackLastCommand()
+    protected virtual void RollbackLastCommand(IReadOnlyModificationCommand modificationCommand)
     {
         _modificationCommands.RemoveAt(_modificationCommands.Count - 1);
 
@@ -154,6 +151,12 @@ protected virtual void RollbackLastCommand()
             RelationalCommandBuilder.RemoveParameterAt(parameterIndex);
             ParameterValues.Remove(parameter.InvariantName);
         }
+
+        // The command's column modifications had their parameter names generated, that needs to be rolled back as well.
+        foreach (var columnModification in modificationCommand.ColumnModifications)
+        {
+            columnModification.ResetParameterNames();
+        }
     }
 
     /// <summary>

diff --git a/src/EFCore.SqlServer/Update/Internal/SqlServerModificationCommandBatch.cs b/src/EFCore.SqlServer/Update/Internal/SqlServerModificationCommandBatch.cs
@@ -53,14 +53,14 @@ public SqlServerModificationCommandBatch(
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    protected override void RollbackLastCommand()
+    protected override void RollbackLastCommand(IReadOnlyModificationCommand modificationCommand)
     {
         if (_pendingBulkInsertCommands.Count > 0)
         {
             _pendingBulkInsertCommands.RemoveAt(_pendingBulkInsertCommands.Count - 1);
         }
 
-        base.RollbackLastCommand();
+        base.RollbackLastCommand(modificationCommand);
     }
 
     /// <summary>

diff --git a/test/EFCore.Relational.Tests/Update/ReaderModificationCommandBatchTest.cs b/test/EFCore.Relational.Tests/Update/ReaderModificationCommandBatchTest.cs
@@ -15,7 +15,7 @@ namespace Microsoft.EntityFrameworkCore.Update;
 public class ReaderModificationCommandBatchTest
 {
     [ConditionalFact]
-    public void AddCommand_adds_command_if_batch_is_valid()
+    public void TryAddCommand_adds_command_if_batch_is_valid()
     {
         var parameterNameGenerator = new ParameterNameGenerator();
 
@@ -73,7 +73,7 @@ public void AddCommand_adds_command_if_batch_is_valid()
     }
 
     [ConditionalFact]
-    public void AddCommand_does_not_add_command_batch_is_invalid()
+    public void TryAddCommand_does_not_add_command_batch_is_invalid()
     {
         var parameterNameGenerator = new ParameterNameGenerator();
 
@@ -131,7 +131,42 @@ public void AddCommand_does_not_add_command_batch_is_invalid()
     }
 
     [ConditionalFact]
-    public void UpdateCommandText_compiles_inserts()
+    public void Parameters_are_properly_managed_when_command_adding_fails()
+    {
+        var entry1 = CreateEntry(EntityState.Added);
+        var command1 = CreateModificationCommand(entry1, new ParameterNameGenerator().GenerateNext, true, null);
+        command1.AddEntry(entry1, true);
+
+        var entry2 = CreateEntry(EntityState.Added);
+        var command2 = CreateModificationCommand(entry2, new ParameterNameGenerator().GenerateNext, true, null);
+        command2.AddEntry(entry2, true);
+
+        var batch1 = new ModificationCommandBatchFake(maxBatchSize: 1);
+        Assert.True(batch1.TryAddCommand(command1));
+        Assert.False(batch1.TryAddCommand(command2));
+        batch1.Complete(moreBatchesExpected: false);
+
+        var batch2 = new ModificationCommandBatchFake(maxBatchSize: 1);
+        Assert.True(batch2.TryAddCommand(command1));
+        batch2.Complete(moreBatchesExpected: false);
+
+        Assert.Equal(2, batch1.StoreCommand.RelationalCommand.Parameters.Count);
+        Assert.Equal("p0", batch1.StoreCommand.RelationalCommand.Parameters[0].InvariantName);
+        Assert.Equal("p1", batch1.StoreCommand.RelationalCommand.Parameters[1].InvariantName);
+
+        Assert.Equal(2, batch2.StoreCommand.RelationalCommand.Parameters.Count);
+        Assert.Equal("p0", batch2.StoreCommand.RelationalCommand.Parameters[0].InvariantName);
+        Assert.Equal("p1", batch2.StoreCommand.RelationalCommand.Parameters[1].InvariantName);
+
+        Assert.Equal(1, batch1.FakeSqlGenerator.AppendBatchHeaderCalls);
+        Assert.Equal(1, batch1.FakeSqlGenerator.AppendInsertOperationCalls);
+
+        Assert.Equal(1, batch2.FakeSqlGenerator.AppendBatchHeaderCalls);
+        Assert.Equal(1, batch2.FakeSqlGenerator.AppendInsertOperationCalls);
+    }
+
+    [ConditionalFact]
+    public void TryAddCommand_with_insert()
     {
         var entry = CreateEntry(EntityState.Added);
 
@@ -147,7 +182,7 @@ public void UpdateCommandText_compiles_inserts()
     }
 
     [ConditionalFact]
-    public void UpdateCommandText_compiles_updates()
+    public void TryAddCommand_with_update()
     {
         var entry = CreateEntry(EntityState.Modified, generateKeyValues: true);
 
@@ -163,7 +198,7 @@ public void UpdateCommandText_compiles_updates()
     }
 
     [ConditionalFact]
-    public void UpdateCommandText_compiles_deletes()
+    public void TryAddCommand_with_delete()
     {
         var entry = CreateEntry(EntityState.Deleted);
 
@@ -179,7 +214,7 @@ public void UpdateCommandText_compiles_deletes()
     }
 
     [ConditionalFact]
-    public void UpdateCommandText_compiles_multiple_commands()
+    public void TryAddCommand_twice()
     {
         var entry = CreateEntry(EntityState.Added);
 
@@ -674,8 +709,8 @@ private class ModificationCommandBatchFake : AffectedCountModificationCommandBat
     {
         private readonly FakeSqlGenerator _fakeSqlGenerator;
 
-        public ModificationCommandBatchFake(IUpdateSqlGenerator sqlGenerator = null)
-            : base(CreateDependencies(sqlGenerator))
+        public ModificationCommandBatchFake(IUpdateSqlGenerator sqlGenerator = null, int? maxBatchSize = null)
+            : base(CreateDependencies(sqlGenerator), maxBatchSize)
         {
             ShouldBeValid = true;