Fix bug where a non-sproc command comes before a sproc command (#29680)

Fixes #29643

Author: roji

src/EFCore.Relational/Update/AffectedCountModificationCommandBatch.cs

@@ -91,12 +91,14 @@ protected override void Consume(RelationalDataReader reader)
                 var parameterCounter = 0;
                 IReadOnlyModificationCommand command;
 
-                for (commandIndex = 0;
-                     commandIndex < ResultSetMappings.Count;
-                     commandIndex++, parameterCounter += command.StoreStoredProcedure!.Parameters.Count)
+                for (commandIndex = 0; commandIndex < ResultSetMappings.Count; commandIndex++, parameterCounter += ParameterCount(command))
                 {
                     command = ModificationCommands[commandIndex];
 
+                    Check.DebugAssert(
+                        command.ColumnModifications.All(c => c.UseParameter),
+                        "This code assumes all column modifications involve a DbParameter (see counting above)");
+
                     if (!ResultSetMappings[commandIndex].HasFlag(ResultSetMapping.HasOutputParameters))
                     {
                         continue;
@@ -210,12 +212,14 @@ protected override async Task ConsumeAsync(
                 var parameterCounter = 0;
                 IReadOnlyModificationCommand command;
 
-                for (commandIndex = 0;
-                     commandIndex < ResultSetMappings.Count;
-                     commandIndex++, parameterCounter += command.StoreStoredProcedure!.Parameters.Count)
+                for (commandIndex = 0; commandIndex < ResultSetMappings.Count; commandIndex++, parameterCounter += ParameterCount(command))
                 {
                     command = ModificationCommands[commandIndex];
 
+                    Check.DebugAssert(
+                        command.ColumnModifications.All(c => c.UseParameter),
+                        "This code assumes all column modifications involve a DbParameter (see counting above)");
+
                     if (!ResultSetMappings[commandIndex].HasFlag(ResultSetMapping.HasOutputParameters))
                     {
                         continue;
@@ -477,6 +481,35 @@ await ThrowAggregateUpdateConcurrencyExceptionAsync(
         return commandIndex - 1;
     }
 
+    private static int ParameterCount(IReadOnlyModificationCommand command)
+    {
+        // As a shortcut, if the command uses a stored procedure, return the number of parameters directly from it.
+        if (command.StoreStoredProcedure is { } storedProcedure)
+        {
+            return storedProcedure.Parameters.Count;
+        }
+
+        // Otherwise we need to count the total parameters used by all column modifications
+        var parameterCount = 0;
+
+        for (var i = 0; i < command.ColumnModifications.Count; i++)
+        {
+            var columnModification = command.ColumnModifications[i];
+
+            if (columnModification.UseCurrentValueParameter)
+            {
+                parameterCount++;
+            }
+
+            if (columnModification.UseOriginalValueParameter)
+            {
+                parameterCount++;
+            }
+        }
+
+        return parameterCount;
+    }
+
     private IReadOnlyList<IUpdateEntry> AggregateEntries(int endIndex, int commandCount)
     {
         var entries = new List<IUpdateEntry>();

test/EFCore.Relational.Specification.Tests/Update/StoredProcedureUpdateTestBase.cs

@@ -1028,6 +1028,48 @@ protected async Task Tpc(bool async, string createSprocSql)
         }
     }
 
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public abstract Task Non_sproc_followed_by_sproc_commands_in_the_same_batch(bool async);
+
+    protected async Task Non_sproc_followed_by_sproc_commands_in_the_same_batch(bool async, string createSprocSql)
+    {
+        var contextFactory = await InitializeAsync<DbContext>(
+            modelBuilder => modelBuilder.Entity<EntityWithAdditionalProperty>()
+                .InsertUsingStoredProcedure(
+                    nameof(EntityWithAdditionalProperty) + "_Insert",
+                    spb => spb
+                        .HasParameter(w => w.Name)
+                        .HasParameter(w => w.Id, pb => pb.IsOutput())
+                        .HasParameter(w => w.AdditionalProperty))
+                .Property(e => e.AdditionalProperty).IsConcurrencyToken(),
+            seed: ctx => CreateStoredProcedures(ctx, createSprocSql));
+
+        await using var context = contextFactory.CreateContext();
+
+        // Prepare by adding an entity
+        var entity1 = new EntityWithAdditionalProperty { Name = "Entity1", AdditionalProperty = 1 };
+        context.Set<EntityWithAdditionalProperty>().Add(entity1);
+
+        using (TestSqlLoggerFactory.SuspendRecordingEvents())
+        {
+            await SaveChanges(context, async);
+        }
+
+        // Now add a second entity and update the first one. The update gets ordered first, and doesn't use a sproc, and then the insertion
+        // does.
+        var entity2 = new EntityWithAdditionalProperty { Name = "Entity2" };
+        context.Set<EntityWithAdditionalProperty>().Add(entity2);
+        entity1.Name = "Entity1_Modified";
+        entity1.AdditionalProperty = 2;
+        await SaveChanges(context, async);
+
+        using (TestSqlLoggerFactory.SuspendRecordingEvents())
+        {
+            Assert.Equal("Entity2", context.Set<EntityWithAdditionalProperty>().Single(b => b.Id == entity2.Id).Name);
+        }
+    }
+
     /// <summary>
     ///     A method to be implement by the provider, to set up a store-generated concurrency token shadow property with the given name.
     /// </summary>

test/EFCore.SqlServer.FunctionalTests/Update/StoredProcedureUpdateSqlServerTest.cs

@@ -665,6 +665,36 @@ AS BEGIN
 """);
     }
 
+    public override async Task Non_sproc_followed_by_sproc_commands_in_the_same_batch(bool async)
+    {
+        await base.Non_sproc_followed_by_sproc_commands_in_the_same_batch(
+            async,
+"""
+CREATE PROCEDURE EntityWithAdditionalProperty_Insert(@Name varchar(max), @Id int OUT, @AdditionalProperty int)
+AS BEGIN
+    INSERT INTO [EntityWithAdditionalProperty] ([Name], [AdditionalProperty]) VALUES (@Name, @AdditionalProperty);
+    SET @Id = SCOPE_IDENTITY();
+END
+""");
+
+        AssertSql(
+"""
+@p2='1'
+@p0='2'
+@p3='1'
+@p1='Entity1_Modified' (Size = 4000)
+@p4='Entity2' (Size = 4000)
+@p5=NULL (Nullable = false) (Direction = Output) (DbType = Int32)
+@p6='0'
+
+SET NOCOUNT ON;
+UPDATE [EntityWithAdditionalProperty] SET [AdditionalProperty] = @p0, [Name] = @p1
+OUTPUT 1
+WHERE [Id] = @p2 AND [AdditionalProperty] = @p3;
+EXEC [EntityWithAdditionalProperty_Insert] @p4, @p5 OUTPUT, @p6;
+""");
+    }
+
     protected override void ConfigureStoreGeneratedConcurrencyToken(EntityTypeBuilder entityTypeBuilder, string propertyName)
         => entityTypeBuilder.Property<byte[]>(propertyName).IsRowVersion();