Query: Fix for query fails with JSON columns and Include collection (#…

…28832)

Co-authored-by: maumar <maumar@microsoft.com>

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

@@ -452,7 +452,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                                 collection: false,
                                 jsonElementParameter,
                                 keyValuesParameter,
-                                outerEntityInstanceParameter: null,
+                                parentEntityExpression: null,
                                 navigation: null);
 
                             var visitedShaperResult = Visit(shaperResult);
@@ -518,7 +518,7 @@ when collectionResultExpression.Navigation is INavigation navigation
                         collection: true,
                         jsonElementParameter,
                         keyValuesParameter,
-                        outerEntityInstanceParameter: null,
+                        parentEntityExpression: null,
                         navigation);
 
                     var visitedShaperResult = Visit(shaperResult);
@@ -786,12 +786,12 @@ when collectionResultExpression.Navigation is INavigation navigation
                                 collection: includeExpression.NavigationExpression is CollectionResultExpression,
                                 jsonElementParameter,
                                 keyValuesParameter,
-                                outerEntityInstanceParameter: (ParameterExpression)entity,
+                                parentEntityExpression: entity,
                                 navigation: (INavigation)includeExpression.Navigation);
 
                             var visitedShaperResult = Visit(shaperResult);
 
-                            _expressions.Add(visitedShaperResult);
+                            _includeExpressions.Add(visitedShaperResult);
 
                             return entity;
                         }
@@ -1074,7 +1074,7 @@ private Expression CreateJsonShapers(
             bool collection,
             ParameterExpression jsonElementParameter,
             ParameterExpression keyValuesParameter,
-            ParameterExpression? outerEntityInstanceParameter,
+            Expression? parentEntityExpression,
             INavigation? navigation)
         {
             var jsonElementShaperLambdaParameter = Expression.Parameter(typeof(JsonElement));
@@ -1141,7 +1141,7 @@ private Expression CreateJsonShapers(
                 keyValuesShaperLambdaParameter,
                 jsonElementShaperLambdaParameter);
 
-            if (outerEntityInstanceParameter != null)
+            if (parentEntityExpression != null)
             {
                 Debug.Assert(navigation != null, "Navigation shouldn't be null when including.");
 
@@ -1152,9 +1152,9 @@ private Expression CreateJsonShapers(
                     navigation.Inverse);
 
                 // inheritance scenario - navigation defined on derived
-                var outerEntityInstanceExpression = outerEntityInstanceParameter.Type != navigation.DeclaringEntityType.ClrType
-                    ? Expression.Convert(outerEntityInstanceParameter, navigation.DeclaringEntityType.ClrType)
-                    : (Expression)outerEntityInstanceParameter;
+                var includingEntityExpression = parentEntityExpression.Type != navigation.DeclaringEntityType.ClrType
+                    ? Expression.Convert(parentEntityExpression, navigation.DeclaringEntityType.ClrType)
+                    : parentEntityExpression;
 
                 if (navigation.IsCollection)
                 {
@@ -1166,15 +1166,15 @@ private Expression CreateJsonShapers(
                             QueryCompilationContext.QueryContextParameter,
                             jsonElementParameter,
                             keyValuesParameter,
-                            outerEntityInstanceExpression,
+                            includingEntityExpression,
                             shaperLambda,
                             fixup);
 
-                    return navigation.DeclaringEntityType.ClrType.IsAssignableFrom(outerEntityInstanceParameter.Type)
+                    return navigation.DeclaringEntityType.ClrType.IsAssignableFrom(parentEntityExpression.Type)
                         ? includeJsonEntityCollectionMethodCall
                         : Expression.IfThen(
                             Expression.TypeIs(
-                                outerEntityInstanceParameter,
+                                parentEntityExpression,
                                 navigation.DeclaringEntityType.ClrType),
                             includeJsonEntityCollectionMethodCall);
                 }
@@ -1187,15 +1187,15 @@ private Expression CreateJsonShapers(
                         QueryCompilationContext.QueryContextParameter,
                         jsonElementParameter,
                         keyValuesParameter,
-                        outerEntityInstanceExpression,
+                        includingEntityExpression,
                         shaperLambda,
                         fixup);
 
-                return navigation.DeclaringEntityType.ClrType.IsAssignableFrom(outerEntityInstanceParameter.Type)
+                return navigation.DeclaringEntityType.ClrType.IsAssignableFrom(parentEntityExpression.Type)
                     ? includeJsonEntityReferenceMethodCall
                     : Expression.IfThen(
                         Expression.TypeIs(
-                            outerEntityInstanceParameter,
+                            parentEntityExpression,
                             navigation.DeclaringEntityType.ClrType),
                         includeJsonEntityReferenceMethodCall);
             }

test/EFCore.Relational.Specification.Tests/Query/JsonQueryFixtureBase.cs

@@ -25,6 +25,8 @@ public virtual ISetSource GetExpectedData()
     public IReadOnlyDictionary<Type, object> EntitySorters { get; } = new Dictionary<Type, Func<object, object>>
     {
         { typeof(JsonEntityBasic), e => ((JsonEntityBasic)e)?.Id },
+        { typeof(JsonEntityBasicForReference), e => ((JsonEntityBasicForReference)e)?.Id },
+        { typeof(JsonEntityBasicForCollection), e => ((JsonEntityBasicForCollection)e)?.Id },
         { typeof(JsonEntityCustomNaming), e => ((JsonEntityCustomNaming)e)?.Id },
         { typeof(JsonEntitySingleOwned), e => ((JsonEntitySingleOwned)e)?.Id },
         { typeof(JsonEntityInheritanceBase), e => ((JsonEntityInheritanceBase)e)?.Id },
@@ -55,6 +57,36 @@ public virtual ISetSource GetExpectedData()
                 }
             }
         },
+        {
+            typeof(JsonEntityBasicForReference), (e, a) =>
+            {
+                Assert.Equal(e == null, a == null);
+                if (a != null)
+                {
+                    var ee = (JsonEntityBasicForReference)e;
+                    var aa = (JsonEntityBasicForReference)a;
+
+                    Assert.Equal(ee.Id, aa.Id);
+                    Assert.Equal(ee.Name, aa.Name);
+                    Assert.Equal(ee.ParentId, aa.ParentId);
+                }
+            }
+        },
+        {
+            typeof(JsonEntityBasicForCollection), (e, a) =>
+            {
+                Assert.Equal(e == null, a == null);
+                if (a != null)
+                {
+                    var ee = (JsonEntityBasicForCollection)e;
+                    var aa = (JsonEntityBasicForCollection)a;
+
+                    Assert.Equal(ee.Id, aa.Id);
+                    Assert.Equal(ee.Name, aa.Name);
+                    Assert.Equal(ee.ParentId, aa.ParentId);
+                }
+            }
+        },
         {
             typeof(JsonOwnedRoot), (e, a) =>
             {
@@ -282,6 +314,8 @@ protected override void Seed(JsonQueryContext context)
     protected override void OnModelCreating(ModelBuilder modelBuilder, DbContext context)
     {
         modelBuilder.Entity<JsonEntityBasic>().Property(x => x.Id).ValueGeneratedNever();
+        modelBuilder.Entity<JsonEntityBasicForReference>().Property(x => x.Id).ValueGeneratedNever();
+        modelBuilder.Entity<JsonEntityBasicForCollection>().Property(x => x.Id).ValueGeneratedNever();
         modelBuilder.Entity<JsonEntityBasic>().OwnsOne(x => x.OwnedReferenceRoot, b =>
         {
             b.ToJson();

test/EFCore.Relational.Specification.Tests/Query/JsonQueryTestBase.cs

@@ -634,5 +634,45 @@ public virtual Task Group_by_on_json_scalar(bool async)
             ss => ss.Set<JsonEntityBasic>()
             .GroupBy(x => x.OwnedReferenceRoot.Name).Select(x => new { x.Key, Count = x.Count() }));
 
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Json_with_include_on_json_entity(bool async)
+        => AssertQuery(
+            async,
+            ss => ss.Set<JsonEntityBasic>().Include(x => x.OwnedReferenceRoot),
+            entryCount: 40);
+
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Json_with_include_on_entity_reference(bool async)
+        => AssertQuery(
+            async,
+            ss => ss.Set<JsonEntityBasic>().Include(x => x.EntityReference),
+            elementAsserter: (e, a) => AssertInclude(
+                e, a,
+                new ExpectedInclude<JsonEntityBasic>(x => x.EntityReference)),
+            entryCount: 41);
 
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Json_with_include_on_entity_collection(bool async)
+        => AssertQuery(
+            async,
+            ss => ss.Set<JsonEntityBasic>().Include(x => x.EntityCollection),
+            elementAsserter: (e, a) => AssertInclude(
+                e, a,
+                new ExpectedInclude<JsonEntityBasic>(x => x.EntityCollection)),
+            entryCount: 43);
+
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Json_with_include_on_entity_collection_and_reference(bool async)
+        => AssertQuery(
+            async,
+            ss => ss.Set<JsonEntityBasic>().Include(x => x.EntityReference).Include(x => x.EntityCollection),
+            elementAsserter: (e, a) => AssertInclude(
+                e, a,
+                new ExpectedInclude<JsonEntityBasic>(x => x.EntityReference),
+                new ExpectedInclude<JsonEntityBasic>(x => x.EntityCollection)),
+            entryCount: 44);
 }

test/EFCore.Relational.Specification.Tests/TestModels/JsonQuery/JsonEntityBasic.cs

@@ -10,5 +10,8 @@ public class JsonEntityBasic
 
         public JsonOwnedRoot OwnedReferenceRoot { get; set; }
         public List<JsonOwnedRoot> OwnedCollectionRoot { get; set; }
+
+        public JsonEntityBasicForReference EntityReference { get; set; }
+        public List<JsonEntityBasicForCollection> EntityCollection { get; set; }
     }
 }

test/EFCore.Relational.Specification.Tests/TestModels/JsonQuery/JsonEntityBasicForCollection.cs

@@ -0,0 +1,14 @@
+// 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.TestModels.JsonQuery
+{
+    public class JsonEntityBasicForCollection
+    {
+        public int Id { get; set; }
+        public string Name { get; set; }
+
+        public int? ParentId { get; set; }
+        public JsonEntityBasic Parent { get; set; }
+    }
+}

test/EFCore.Relational.Specification.Tests/TestModels/JsonQuery/JsonEntityBasicForReference.cs

@@ -0,0 +1,14 @@
+// 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.TestModels.JsonQuery
+{
+    public class JsonEntityBasicForReference
+    {
+        public int Id { get; set; }
+        public string Name { get; set; }
+
+        public int? ParentId { get; set; }
+        public JsonEntityBasic Parent { get; set; }
+    }
+}

test/EFCore.Relational.Specification.Tests/TestModels/JsonQuery/JsonQueryContext.cs

@@ -11,18 +11,26 @@ public JsonQueryContext(DbContextOptions options)
         }
 
         public DbSet<JsonEntityBasic> JsonEntitiesBasic { get; set; }
+        public DbSet<JsonEntityBasicForReference> JsonEntitiesBasicForReference { get; set; }
+        public DbSet<JsonEntityBasicForCollection> JsonEntitiesBasicForCollection { get; set; }
         public DbSet<JsonEntityCustomNaming> JsonEntitiesCustomNaming { get; set; }
         public DbSet<JsonEntitySingleOwned> JsonEntitiesSingleOwned { get; set; }
         public DbSet<JsonEntityInheritanceBase> JsonEntitiesInheritance { get; set; }
 
         public static void Seed(JsonQueryContext context)
         {
             var jsonEntitiesBasic = JsonQueryData.CreateJsonEntitiesBasic();
+            var jsonEntitiesBasicForReference = JsonQueryData.CreateJsonEntitiesBasicForReference();
+            var jsonEntitiesBasicForCollection = JsonQueryData.CreateJsonEntitiesBasicForCollection();
+            JsonQueryData.WireUp(jsonEntitiesBasic, jsonEntitiesBasicForReference, jsonEntitiesBasicForCollection);
+
             var jsonEntitiesCustomNaming = JsonQueryData.CreateJsonEntitiesCustomNaming();
             var jsonEntitiesSingleOwned = JsonQueryData.CreateJsonEntitiesSingleOwned();
             var jsonEntitiesInheritance = JsonQueryData.CreateJsonEntitiesInheritance();
 
             context.JsonEntitiesBasic.AddRange(jsonEntitiesBasic);
+            context.JsonEntitiesBasicForReference.AddRange(jsonEntitiesBasicForReference);
+            context.JsonEntitiesBasicForCollection.AddRange(jsonEntitiesBasicForCollection);
             context.JsonEntitiesCustomNaming.AddRange(jsonEntitiesCustomNaming);
             context.JsonEntitiesSingleOwned.AddRange(jsonEntitiesSingleOwned);
             context.JsonEntitiesInheritance.AddRange(jsonEntitiesInheritance);

test/EFCore.Relational.Specification.Tests/TestModels/JsonQuery/JsonQueryData.cs

@@ -8,12 +8,18 @@ public class JsonQueryData : ISetSource
     public JsonQueryData()
     {
         JsonEntitiesBasic = CreateJsonEntitiesBasic();
+        JsonEntitiesBasicForReference = CreateJsonEntitiesBasicForReference();
+        JsonEntitiesBasicForCollection = CreateJsonEntitiesBasicForCollection();
+        WireUp(JsonEntitiesBasic, JsonEntitiesBasicForReference, JsonEntitiesBasicForCollection);
+
         JsonEntitiesCustomNaming = CreateJsonEntitiesCustomNaming();
         JsonEntitiesSingleOwned = CreateJsonEntitiesSingleOwned();
         JsonEntitiesInheritance = CreateJsonEntitiesInheritance();
     }
 
     public IReadOnlyList<JsonEntityBasic> JsonEntitiesBasic { get; }
+    public IReadOnlyList<JsonEntityBasicForReference> JsonEntitiesBasicForReference { get; }
+    public IReadOnlyList<JsonEntityBasicForCollection> JsonEntitiesBasicForCollection { get; }
     public IReadOnlyList<JsonEntityCustomNaming> JsonEntitiesCustomNaming { get; set; }
     public IReadOnlyList<JsonEntitySingleOwned> JsonEntitiesSingleOwned { get; set; }
     public IReadOnlyList<JsonEntityInheritanceBase> JsonEntitiesInheritance { get; set; }
@@ -277,6 +283,46 @@ public static IReadOnlyList<JsonEntityBasic> CreateJsonEntitiesBasic()
         return new List<JsonEntityBasic> { entity1 };
     }
 
+
+    public static IReadOnlyList<JsonEntityBasicForReference> CreateJsonEntitiesBasicForReference()
+    {
+        var entity1 = new JsonEntityBasicForReference { Id = 1, Name = "EntityReference1" };
+
+        return new List<JsonEntityBasicForReference> { entity1 };
+    }
+
+    public static IReadOnlyList<JsonEntityBasicForCollection> CreateJsonEntitiesBasicForCollection()
+    {
+        var entity1 = new JsonEntityBasicForCollection { Id = 1, Name = "EntityCollection1" };
+        var entity2 = new JsonEntityBasicForCollection { Id = 2, Name = "EntityCollection2" };
+        var entity3 = new JsonEntityBasicForCollection { Id = 3, Name = "EntityCollection3" };
+
+        return new List<JsonEntityBasicForCollection> { entity1, entity2, entity3 };
+    }
+
+    public static void WireUp(
+        IReadOnlyList<JsonEntityBasic> entitiesBasic,
+        IReadOnlyList<JsonEntityBasicForReference> entitiesBasicForReference,
+        IReadOnlyList<JsonEntityBasicForCollection> entitiesBasicForCollection)
+    {
+        entitiesBasic[0].EntityReference = entitiesBasicForReference[0];
+        entitiesBasicForReference[0].Parent = entitiesBasic[0];
+        entitiesBasicForReference[0].ParentId = entitiesBasic[0].Id;
+
+        entitiesBasic[0].EntityCollection = new List<JsonEntityBasicForCollection> {
+            entitiesBasicForCollection[0],
+            entitiesBasicForCollection[1],
+            entitiesBasicForCollection[2]
+        };
+
+        entitiesBasicForCollection[0].Parent = entitiesBasic[0];
+        entitiesBasicForCollection[0].ParentId = entitiesBasic[0].Id;
+        entitiesBasicForCollection[1].Parent = entitiesBasic[0];
+        entitiesBasicForCollection[1].ParentId = entitiesBasic[0].Id;
+        entitiesBasicForCollection[2].Parent = entitiesBasic[0];
+        entitiesBasicForCollection[2].ParentId = entitiesBasic[0].Id;
+    }
+
     public static IReadOnlyList<JsonEntityCustomNaming> CreateJsonEntitiesCustomNaming()
     {
         var e1_r_r = new JsonOwnedCustomNameBranch

test/EFCore.SqlServer.FunctionalTests/Query/JsonQuerySqlServerTest.cs

@@ -553,6 +553,48 @@ FROM [JsonEntitiesBasic] AS [j]
 GROUP BY [t].[Key]");
     }
 
+    public override async Task Json_with_include_on_json_entity(bool async)
+    {
+        await base.Json_with_include_on_json_entity(async);
+
+        AssertSql(
+            @"SELECT [j].[Id], [j].[Name], JSON_QUERY([j].[OwnedCollectionRoot],'$'), JSON_QUERY([j].[OwnedReferenceRoot],'$')
+FROM [JsonEntitiesBasic] AS [j]");
+    }
+
+    public override async Task Json_with_include_on_entity_reference(bool async)
+    {
+        await base.Json_with_include_on_entity_reference(async);
+
+        AssertSql(
+            @"SELECT [j].[Id], [j].[Name], JSON_QUERY([j].[OwnedCollectionRoot],'$'), JSON_QUERY([j].[OwnedReferenceRoot],'$'), [j0].[Id], [j0].[Name], [j0].[ParentId]
+FROM [JsonEntitiesBasic] AS [j]
+LEFT JOIN [JsonEntitiesBasicForReference] AS [j0] ON [j].[Id] = [j0].[ParentId]");
+    }
+
+    public override async Task Json_with_include_on_entity_collection(bool async)
+    {
+        await base.Json_with_include_on_entity_collection(async);
+
+        AssertSql(
+            @"SELECT [j].[Id], [j].[Name], JSON_QUERY([j].[OwnedCollectionRoot],'$'), JSON_QUERY([j].[OwnedReferenceRoot],'$'), [j0].[Id], [j0].[Name], [j0].[ParentId]
+FROM [JsonEntitiesBasic] AS [j]
+LEFT JOIN [JsonEntitiesBasicForCollection] AS [j0] ON [j].[Id] = [j0].[ParentId]
+ORDER BY [j].[Id]");
+    }
+
+    public override async Task Json_with_include_on_entity_collection_and_reference(bool async)
+    {
+        await base.Json_with_include_on_entity_collection_and_reference(async);
+
+        AssertSql(
+            @"SELECT [j].[Id], [j].[Name], JSON_QUERY([j].[OwnedCollectionRoot],'$'), JSON_QUERY([j].[OwnedReferenceRoot],'$'), [j0].[Id], [j0].[Name], [j0].[ParentId], [j1].[Id], [j1].[Name], [j1].[ParentId]
+FROM [JsonEntitiesBasic] AS [j]
+LEFT JOIN [JsonEntitiesBasicForReference] AS [j0] ON [j].[Id] = [j0].[ParentId]
+LEFT JOIN [JsonEntitiesBasicForCollection] AS [j1] ON [j].[Id] = [j1].[ParentId]
+ORDER BY [j].[Id], [j0].[Id]");
+    }
+
     private void AssertSql(params string[] expected)
         => Fixture.TestSqlLoggerFactory.AssertBaseline(expected);
 }