Inject JSON shaper code before the entity gets tracked (#36446)

Closes #36433

Author: roji

src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs

@@ -484,6 +484,27 @@ private static Expression CreateReaderColumnsExpression(
         return result;
     }
 
+    /// <summary>
+    ///     Called after a structural type is materialized, but before it's handed off to the change tracker.
+    ///     Here we inject the JSON shapers for any complex JSON properties the type has.
+    /// </summary>
+    /// <remarks>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     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.
+    /// </remarks>
+    public override void AddStructuralTypeInitialization(
+        StructuralTypeShaperExpression shaper,
+        ParameterExpression instanceVariable,
+        List<ParameterExpression> variables,
+        List<Expression> expressions)
+    {
+        Check.DebugAssert(_currentShaperProcessor is not null);
+
+        _currentShaperProcessor.ProcessTopLevelComplexJsonProperties(shaper, instanceVariable, expressions);
+    }
+
     private Expression CreateRelationalCommandResolverExpression(Expression queryExpression)
     {
         // In the regular case, we generate code that accesses the RelationalCommandCache (which invokes the 2nd part of the

src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.cs

@@ -140,7 +140,7 @@ IServiceProperty serviceProperty
                 => serviceProperty.ParameterBinding.BindToParameter(bindingInfo),
 
             IComplexProperty { IsCollection: true } complexProperty
-                => Expression.Default(complexProperty.ClrType), // Initialize collections to null, they'll be populated separately
+                => Default(complexProperty.ClrType), // Initialize collections to null, they'll be populated separately
 
             IComplexProperty complexProperty
                 => CreateMaterializeExpression(

src/EFCore/Query/Internal/StructuralTypeMaterializerSource.cs

@@ -150,7 +150,7 @@ public static Expression BuildMemberAccessForEntityOrComplexType(
     ///     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>
-    public static Expression<Func<MaterializerLiftableConstantContext, object>> BuildMemberAccessLambdaForEntityOrComplexType(
+    public static Expression<Func<MaterializerLiftableConstantContext, object>> BuildMemberAccessLambdaForStructuralType(
         ITypeBase type)
     {
         var prm = Parameter(typeof(MaterializerLiftableConstantContext));

src/EFCore/Query/LiftableConstantExpressionHelpers.cs

@@ -55,6 +55,7 @@ protected ShapedQueryCompilingExpressionVisitor(
 
         _structuralTypeMaterializerInjector =
             new StructuralTypeMaterializerInjector(
+                this,
                 dependencies.EntityMaterializerSource,
                 dependencies.LiftableConstantFactory,
                 queryCompilationContext.QueryTrackingBehavior,
@@ -233,12 +234,12 @@ protected override Expression VisitConstant(ConstantExpression constantExpressio
             {
                 { Value: IEntityType entityTypeValue } => liftableConstantFactory.CreateLiftableConstant(
                     constantExpression.Value,
-                    LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForEntityOrComplexType(entityTypeValue),
+                    LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForStructuralType(entityTypeValue),
                     entityTypeValue.ShortName() + "EntityType",
                     constantExpression.Type),
                 { Value: IComplexType complexTypeValue } => liftableConstantFactory.CreateLiftableConstant(
                     constantExpression.Value,
-                    LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForEntityOrComplexType(complexTypeValue),
+                    LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForStructuralType(complexTypeValue),
                     complexTypeValue.ShortName() + "ComplexType",
                     constantExpression.Type),
                 { Value: IProperty propertyValue } => liftableConstantFactory.CreateLiftableConstant(
@@ -361,8 +362,20 @@ protected override Expression VisitExtension(Expression extensionExpression)
         }
     }
 
+    /// <summary>
+    ///     Called after a structural type is materialized, but before it's handed off to the change tracker.
+    /// </summary>
+    public virtual void AddStructuralTypeInitialization(
+        StructuralTypeShaperExpression shaper,
+        ParameterExpression instanceVariable,
+        List<ParameterExpression> variables,
+        List<Expression> expressions)
+    {
+    }
+
     private sealed class StructuralTypeMaterializerInjector(
-        IStructuralTypeMaterializerSource entityMaterializerSource,
+        ShapedQueryCompilingExpressionVisitor shapedQueryCompiler,
+        IStructuralTypeMaterializerSource materializerSource,
         ILiftableConstantFactory liftableConstantFactory,
         QueryTrackingBehavior queryTrackingBehavior,
         bool supportsPrecompiledQuery)
@@ -426,10 +439,10 @@ bool ContainsOwner(IEntityType? owner)
 
         protected override Expression VisitExtension(Expression extensionExpression)
             => extensionExpression is StructuralTypeShaperExpression shaper
-                ? ProcessEntityShaper(shaper)
+                ? ProcessStructuralTypeShaper(shaper)
                 : base.VisitExtension(extensionExpression);
 
-        private Expression ProcessEntityShaper(StructuralTypeShaperExpression shaper)
+        private Expression ProcessStructuralTypeShaper(StructuralTypeShaperExpression shaper)
         {
             _currentEntityIndex++;
 
@@ -565,7 +578,7 @@ private Expression ProcessEntityShaper(StructuralTypeShaperExpression shaper)
                                         supportsPrecompiledQuery
                                             ? liftableConstantFactory.CreateLiftableConstant(
                                                 typeBase,
-                                                LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForEntityOrComplexType(typeBase),
+                                                LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForStructuralType(typeBase),
                                                 typeBase.Name + "EntityType",
                                                 typeof(IEntityType))
                                             : Constant(typeBase),
@@ -606,7 +619,7 @@ private Expression MaterializeEntity(
             ParameterExpression instanceVariable,
             ParameterExpression? entryVariable)
         {
-            var typeBase = shaper.StructuralType;
+            var structuralType = shaper.StructuralType;
 
             var expressions = new List<Expression>();
             var variables = new List<ParameterExpression>();
@@ -626,7 +639,7 @@ private Expression MaterializeEntity(
                             typeof(ISnapshot))
                         : Constant(Snapshot.Empty, typeof(ISnapshot))));
 
-            var returnType = typeBase.ClrType;
+            var returnType = structuralType.ClrType;
             var valueBufferExpression = Call(materializationContextVariable, MaterializationContext.GetValueBufferMethod);
 
             var materializationConditionBody = ReplacingExpressionVisitor.Replace(
@@ -637,23 +650,23 @@ private Expression MaterializeEntity(
             var expressionContext = (returnType, materializationContextVariable, concreteEntityTypeVariable, shadowValuesVariable);
             expressions.Add(Assign(concreteEntityTypeVariable, materializationConditionBody));
 
-            var (primaryKey, concreteEntityTypes) = typeBase is IEntityType entityType
+            var (primaryKey, concreteStructuralTypes) = structuralType is IEntityType entityType
                 ? (entityType.FindPrimaryKey(), entityType.GetConcreteDerivedTypesInclusive().Cast<ITypeBase>().ToArray())
-                : (null, [typeBase]);
+                : (null, [structuralType]);
 
-            var switchCases = new SwitchCase[concreteEntityTypes.Length];
-            for (var i = 0; i < concreteEntityTypes.Length; i++)
+            var switchCases = new SwitchCase[concreteStructuralTypes.Length];
+            for (var i = 0; i < concreteStructuralTypes.Length; i++)
             {
-                var concreteEntityType = concreteEntityTypes[i];
+                var concreteStructuralType = concreteStructuralTypes[i];
                 switchCases[i] = SwitchCase(
-                    CreateFullMaterializeExpression(concreteEntityTypes[i], expressionContext),
+                    CreateFullMaterializeExpression(concreteStructuralTypes[i], expressionContext),
                     supportsPrecompiledQuery
                         ? liftableConstantFactory.CreateLiftableConstant(
-                            concreteEntityTypes[i],
-                            LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForEntityOrComplexType(concreteEntityType),
-                            concreteEntityType.ShortName() + (typeBase is IEntityType ? "EntityType" : "ComplexType"),
-                            typeBase is IEntityType ? typeof(IEntityType) : typeof(IComplexType))
-                        : Constant(concreteEntityTypes[i], typeBase is IEntityType ? typeof(IEntityType) : typeof(IComplexType)));
+                            concreteStructuralTypes[i],
+                            LiftableConstantExpressionHelpers.BuildMemberAccessLambdaForStructuralType(concreteStructuralType),
+                            concreteStructuralType.ShortName() + (structuralType is IEntityType ? "EntityType" : "ComplexType"),
+                            structuralType is IEntityType ? typeof(IEntityType) : typeof(IComplexType))
+                        : Constant(concreteStructuralTypes[i], structuralType is IEntityType ? typeof(IEntityType) : typeof(IComplexType)));
             }
 
             var materializationExpression = Switch(
@@ -663,9 +676,11 @@ private Expression MaterializeEntity(
 
             expressions.Add(Assign(instanceVariable, materializationExpression));
 
+            shapedQueryCompiler.AddStructuralTypeInitialization(shaper, instanceVariable, variables, expressions);
+
             if (_queryStateManager && primaryKey is not null)
             {
-                if (typeBase is IEntityType entityType2)
+                if (structuralType is IEntityType entityType2)
                 {
                     foreach (var et in entityType2.GetAllBaseTypes().Concat(entityType2.GetDerivedTypesInclusive()))
                     {
@@ -696,7 +711,7 @@ private Expression MaterializeEntity(
         }
 
         private BlockExpression CreateFullMaterializeExpression(
-            ITypeBase concreteTypeBase,
+            ITypeBase concreteStructuralType,
             (Type ReturnType,
                 ParameterExpression MaterializationContextVariable,
                 ParameterExpression ConcreteEntityTypeVariable,
@@ -709,14 +724,14 @@ private BlockExpression CreateFullMaterializeExpression(
 
             var blockExpressions = new List<Expression>(2);
 
-            var materializer = entityMaterializerSource
+            var materializer = materializerSource
                 .CreateMaterializeExpression(
                     new StructuralTypeMaterializerSourceParameters(
-                        concreteTypeBase, "instance", queryTrackingBehavior), materializationContextVariable);
+                        concreteStructuralType, "instance", queryTrackingBehavior), materializationContextVariable);
 
             // TODO: Properly support shadow properties for complex types #35613
             if (_queryStateManager
-                && concreteTypeBase is IRuntimeEntityType { ShadowPropertyCount: > 0 } runtimeEntityType)
+                && concreteStructuralType is IRuntimeEntityType { ShadowPropertyCount: > 0 } runtimeEntityType)
             {
                 var valueBufferExpression = Call(
                     materializationContextVariable, MaterializationContext.GetValueBufferMethod);

src/EFCore/Query/ShapedQueryCompilingExpressionVisitor.cs

@@ -11,7 +11,7 @@ public abstract class ComplexCollectionJsonUpdateTestBase<TFixture>(TFixture fix
     protected ComplexCollectionJsonContext CreateContext()
         => (ComplexCollectionJsonContext)Fixture.CreateContext();
 
-    [ConditionalFact(Skip = "Issue #36433")]
+    [ConditionalFact]
     public virtual Task Add_element_to_complex_collection_mapped_to_json()
         => TestHelpers.ExecuteWithStrategyInTransactionAsync(
             CreateContext,
@@ -26,17 +26,10 @@ public virtual Task Add_element_to_complex_collection_mapped_to_json()
 
                 company.Contacts!.Add(new Contact { Name = "New Contact", PhoneNumbers = ["555-0000"] });
 
-                Assert.Equal("""
-CompanyWithComplexCollections {Id: 1} Unchanged
-    Id: 1 PK
-    Name: 'Test Company'
-    Contacts (Complex: List<Contact>)
-    Department (Complex: Department)
-      Budget: 10000.00
-      Name: 'Initial Department'
-    Employees (Complex: List<Employee>)
-
-""", context.ChangeTracker.DebugView.LongView);
+                Assert.Contains("Contacts (Complex: List<Contact>)", context.ChangeTracker.DebugView.LongView);
+                Assert.Contains("Department (Complex: Department)", context.ChangeTracker.DebugView.LongView);
+                Assert.Contains("Name: 'Initial Department'", context.ChangeTracker.DebugView.LongView);
+                Assert.Contains("Employees (Complex: List<Employee>)", context.ChangeTracker.DebugView.LongView);
 
                 ClearLog();
                 await context.SaveChangesAsync();
@@ -126,7 +119,7 @@ public virtual Task Move_elements_in_complex_collection_mapped_to_json()
                 }
             });
 
-    [ConditionalFact(Skip = "Issue #36433")]
+    [ConditionalFact]
     public virtual Task Change_complex_collection_mapped_to_json_to_null_and_to_empty()
         => TestHelpers.ExecuteWithStrategyInTransactionAsync(
             CreateContext,
@@ -354,7 +347,7 @@ public virtual Task Modify_nested_complex_property_in_complex_collection_mapped_
                 }
             });
 
-    [ConditionalFact(Skip = "Issue #36433")]
+    [ConditionalFact]
     public virtual Task Set_complex_collection_to_null_mapped_to_json()
         => TestHelpers.ExecuteWithStrategyInTransactionAsync(
             CreateContext,
@@ -499,7 +492,7 @@ public virtual Task Complex_collection_with_empty_nested_collections_mapped_to_j
                 }
             });
 
-    [ConditionalFact(Skip = "Issue #36433")]
+    [ConditionalFact]
     public virtual Task Set_complex_property_mapped_to_json_to_null()
         => TestHelpers.ExecuteWithStrategyInTransactionAsync(
             CreateContext,