Fix pruning of UpdateExpression

Fixes dotnet#31407
roji · Dec 25, 2023 · e4a6fff · e4a6fff
1 parent 216c83e
commit e4a6fff
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/src/EFCore.Relational/Query/SqlTreePruner.cs b/src/EFCore.Relational/Query/SqlTreePruner.cs
@@ -65,9 +65,13 @@ protected override Expression VisitExtension(Expression node)
                 return deleteExpression.Update(PruneTopLevelSelect(deleteExpression.SelectExpression));
 
             case UpdateExpression updateExpression:
+                // Note that we must visit the setters before we visit the select, since the setters can reference tables inside it.
+                var visitedSetters = updateExpression.ColumnValueSetters
+                    .Select(e => e with { Value = (SqlExpression)Visit(e.Value) })
+                    .ToList();
                 return updateExpression.Update(
                     PruneTopLevelSelect(updateExpression.SelectExpression),
-                    updateExpression.ColumnValueSetters.Select(e => e with { Value = (SqlExpression)Visit(e.Value) }).ToList());
+                    visitedSetters);
 
             // The following remaining cases deal with recursive visitation (i.e. non-top-level things)
 

diff --git a/test/EFCore.SqlServer.FunctionalTests/BulkUpdates/NonSharedModelBulkUpdatesSqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/BulkUpdates/NonSharedModelBulkUpdatesSqlServerTest.cs
@@ -90,10 +90,19 @@ FROM [Blogs] AS [b]
 """);
     }
 
-    // #31407
-    public override Task Update_non_main_table_in_entity_with_entity_splitting(bool async)
-        => Assert.ThrowsAnyAsync<Exception>(
-            () => base.Update_non_main_table_in_entity_with_entity_splitting(async));
+    public override async Task Update_non_main_table_in_entity_with_entity_splitting(bool async)
+    {
+        await base.Update_non_main_table_in_entity_with_entity_splitting(async);
+
+        AssertSql(
+            """
+UPDATE [b0]
+SET [b0].[Rating] = CAST(LEN([b0].[Title]) AS int),
+    [b0].[Title] = CONVERT(varchar(11), [b0].[Rating])
+FROM [Blogs] AS [b]
+INNER JOIN [BlogsPart1] AS [b0] ON [b].[Id] = [b0].[Id]
+""");
+    }
 
     public override async Task Delete_entity_with_auto_include(bool async)
     {