[rc2] Fix transformation from JsonQueryExpression to OPENJSON

src/EFCore.SqlServer/Query/Internal/SqlServerQueryableMethodTranslatingExpressionVisitor.cs

@@ -22,8 +22,6 @@ public class SqlServerQueryableMethodTranslatingExpressionVisitor : RelationalQu
     private readonly ISqlExpressionFactory _sqlExpressionFactory;
     private readonly ISqlServerSingletonOptions _sqlServerSingletonOptions;
 
-    private RelationalTypeMapping? _nvarcharMaxTypeMapping;
-
     /// <summary>
     ///     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
@@ -239,6 +237,8 @@ protected override Expression VisitExtension(Expression extensionExpression)
     /// </summary>
     protected override ShapedQueryExpression TransformJsonQueryToTable(JsonQueryExpression jsonQueryExpression)
     {
+        var structuralType = jsonQueryExpression.StructuralType;
+
         // Calculate the table alias for the OPENJSON expression based on the last named path segment
         // (or the JSON column name if there are none)
         var lastNamedPathSegment = jsonQueryExpression.Path.LastOrDefault(ps => ps.PropertyName is not null);
@@ -251,7 +251,8 @@ protected override ShapedQueryExpression TransformJsonQueryToTable(JsonQueryExpr
         var columnInfos = new List<SqlServerOpenJsonExpression.ColumnInfo>();
 
         // We're only interested in properties which actually exist in the JSON, filter out uninteresting shadow keys
-        foreach (var property in jsonQueryExpression.StructuralType.GetPropertiesInHierarchy())
+        // (for owned JSON entities)
+        foreach (var property in structuralType.GetPropertiesInHierarchy())
         {
             if (property.GetJsonPropertyName() is { } jsonPropertyName)
             {
@@ -266,46 +267,40 @@ protected override ShapedQueryExpression TransformJsonQueryToTable(JsonQueryExpr
             }
         }
 
-        switch (jsonQueryExpression.StructuralType)
+        // Find the container column in the relational model to get its type mapping
+        // Note that we assume exactly one column with the given name mapped to the entity (despite entity splitting).
+        // See #36647 and #36646 about improving this.
+        var containerColumnName = structuralType.GetContainerColumnName();
+        var containerColumn = structuralType.ContainingEntityType.GetTableMappings()
+            .SelectMany(m => m.Table.Columns)
+            .Where(c => c.Name == containerColumnName)
+            .Single();
+
+        var nestedJsonPropertyNames = jsonQueryExpression.StructuralType switch
         {
-            case IEntityType entityType:
-                // Navigations represent nested JSON owned entities, which we also add to the OPENJSON WITH clause, but with AS JSON.
-                foreach (var navigation in entityType.GetNavigationsInHierarchy()
-                             .Where(n => n.ForeignKey.IsOwnership
-                                 && n.TargetEntityType.IsMappedToJson()
-                                 && n.ForeignKey.PrincipalToDependent == n))
-                {
-                    var jsonPropertyName = navigation.TargetEntityType.GetJsonPropertyName();
-                    Check.DebugAssert(jsonPropertyName is not null, $"No JSON property name for navigation {navigation.Name}");
+            IEntityType entityType
+                => entityType.GetNavigationsInHierarchy()
+                    .Where(n => n.ForeignKey.IsOwnership
+                        && n.TargetEntityType.IsMappedToJson()
+                        && n.ForeignKey.PrincipalToDependent == n)
+                    .Select(n => n.TargetEntityType.GetJsonPropertyName() ?? throw new UnreachableException()),
 
-                    AddStructuralColumnInfo(jsonPropertyName);
-                }
+            IComplexType complexType
+                => complexType.GetComplexProperties().Select(p => p.ComplexType.GetJsonPropertyName() ?? throw new UnreachableException()),
 
-                break;
+            _ => throw new UnreachableException()
+        };
 
-            case IComplexType complexType:
-                foreach (var complexProperty in complexType.GetComplexProperties())
+        foreach (var jsonPropertyName in nestedJsonPropertyNames)
+        {
+            columnInfos.Add(
+                new SqlServerOpenJsonExpression.ColumnInfo
                 {
-                    var jsonPropertyName = complexProperty.ComplexType.GetJsonPropertyName();
-                    Check.DebugAssert(jsonPropertyName is not null, $"No JSON property name for complex property {complexProperty.Name}");
-
-                    AddStructuralColumnInfo(jsonPropertyName);
-                }
-
-                break;
-
-            default:
-                throw new UnreachableException();
-
-                void AddStructuralColumnInfo(string jsonPropertyName)
-                    => columnInfos.Add(
-                        new SqlServerOpenJsonExpression.ColumnInfo
-                        {
-                            Name = jsonPropertyName,
-                            TypeMapping = _nvarcharMaxTypeMapping ??= _typeMappingSource.FindMapping("nvarchar(max)")!,
-                            Path = [new PathSegment(jsonPropertyName)],
-                            AsJson = true
-                        });
+                    Name = jsonPropertyName,
+                    TypeMapping = containerColumn.StoreTypeMapping,
+                    Path = [new PathSegment(jsonPropertyName)],
+                    AsJson = true
+                });
         }
 
         var openJsonExpression = new SqlServerOpenJsonExpression(

test/EFCore.SqlServer.FunctionalTests/Query/Associations/ComplexJson/ComplexJsonCollectionSqlServerTest.cs

@@ -3,6 +3,8 @@
 
 namespace Microsoft.EntityFrameworkCore.Query.Associations.ComplexJson;
 
+using Microsoft.Data.SqlClient;
+
 public class ComplexJsonCollectionSqlServerTest(ComplexJsonSqlServerFixture fixture, ITestOutputHelper testOutputHelper)
     : ComplexJsonCollectionRelationalTestBase<ComplexJsonSqlServerFixture>(fixture, testOutputHelper)
 {
@@ -58,10 +60,37 @@ ORDER BY [r0].[Id]
 
     public override async Task Distinct()
     {
-        await base.Distinct();
+        if (Fixture.UsingJsonType)
+        {
+            // The json data type cannot be selected as DISTINCT because it is not comparable.
+            await Assert.ThrowsAsync<SqlException>(base.Distinct);
 
-        AssertSql(
-            """
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
+FROM [RootEntity] AS [r]
+WHERE (
+    SELECT COUNT(*)
+    FROM (
+        SELECT DISTINCT [r0].[Id], [r0].[Int], [r0].[Name], [r0].[String], [r0].[NestedCollection] AS [c], [r0].[OptionalNested] AS [c0], [r0].[RequiredNested] AS [c1]
+        FROM OPENJSON([r].[RelatedCollection], '$') WITH (
+            [Id] int '$.Id',
+            [Int] int '$.Int',
+            [Name] nvarchar(max) '$.Name',
+            [String] nvarchar(max) '$.String',
+            [NestedCollection] json '$.NestedCollection' AS JSON,
+            [OptionalNested] json '$.OptionalNested' AS JSON,
+            [RequiredNested] json '$.RequiredNested' AS JSON
+        ) AS [r0]
+    ) AS [r1]) = 2
+""");
+        }
+        else
+        {
+            await base.Distinct();
+
+            AssertSql(
+                """
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
 WHERE (
@@ -79,6 +108,7 @@ [RequiredNested] nvarchar(max) '$.RequiredNested' AS JSON
         ) AS [r0]
     ) AS [r1]) = 2
 """);
+        }
     }
 
     public override async Task Distinct_projected(QueryTrackingBehavior queryTrackingBehavior)
@@ -90,10 +120,29 @@ public override async Task Distinct_projected(QueryTrackingBehavior queryTrackin
 
     public override async Task Distinct_over_projected_nested_collection()
     {
-        await base.Distinct_over_projected_nested_collection();
+        if (Fixture.UsingJsonType)
+        {
+            // The json data type cannot be selected as DISTINCT because it is not comparable.
+            await Assert.ThrowsAsync<SqlException>(base.Distinct_over_projected_nested_collection);
 
-        AssertSql(
-            """
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
+FROM [RootEntity] AS [r]
+WHERE (
+    SELECT COUNT(*)
+    FROM (
+        SELECT DISTINCT [r0].[NestedCollection] AS [c]
+        FROM OPENJSON([r].[RelatedCollection], '$') WITH ([NestedCollection] json '$.NestedCollection' AS JSON) AS [r0]
+    ) AS [r1]) = 2
+""");
+        }
+        else
+        {
+            await base.Distinct_over_projected_nested_collection();
+
+            AssertSql(
+                """
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
 WHERE (
@@ -103,6 +152,7 @@ SELECT DISTINCT [r0].[NestedCollection] AS [c]
         FROM OPENJSON([r].[RelatedCollection], '$') WITH ([NestedCollection] nvarchar(max) '$.NestedCollection' AS JSON) AS [r0]
     ) AS [r1]) = 2
 """);
+        }
     }
 
     public override async Task Distinct_over_projected_filtered_nested_collection()
@@ -246,8 +296,24 @@ public override async Task Select_within_Select_within_Select_with_aggregates()
     {
         await base.Select_within_Select_within_Select_with_aggregates();
 
-        AssertSql(
-            """
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+SELECT (
+    SELECT COALESCE(SUM([s].[value]), 0)
+    FROM OPENJSON([r].[RelatedCollection], '$') WITH ([NestedCollection] json '$.NestedCollection' AS JSON) AS [r0]
+    OUTER APPLY (
+        SELECT MAX([n].[Int]) AS [value]
+        FROM OPENJSON([r0].[NestedCollection], '$') WITH ([Int] int '$.Int') AS [n]
+    ) AS [s])
+FROM [RootEntity] AS [r]
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
 SELECT (
     SELECT COALESCE(SUM([s].[value]), 0)
     FROM OPENJSON([r].[RelatedCollection], '$') WITH ([NestedCollection] nvarchar(max) '$.NestedCollection' AS JSON) AS [r0]
@@ -257,6 +323,7 @@ FROM OPENJSON([r0].[NestedCollection], '$') WITH ([Int] int '$.Int') AS [n]
     ) AS [s])
 FROM [RootEntity] AS [r]
 """);
+        }
     }
 
     [ConditionalFact]

test/EFCore.SqlServer.FunctionalTests/Query/Associations/ComplexJson/ComplexJsonProjectionSqlServerTest.cs

@@ -233,8 +233,27 @@ public override async Task SelectMany_related_collection(QueryTrackingBehavior q
     {
         await base.SelectMany_related_collection(queryTrackingBehavior);
 
-        AssertSql(
-            """
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+SELECT [r0].[Id], [r0].[Int], [r0].[Name], [r0].[String], [r0].[NestedCollection], [r0].[OptionalNested], [r0].[RequiredNested]
+FROM [RootEntity] AS [r]
+CROSS APPLY OPENJSON([r].[RelatedCollection], '$') WITH (
+    [Id] int '$.Id',
+    [Int] int '$.Int',
+    [Name] nvarchar(max) '$.Name',
+    [String] nvarchar(max) '$.String',
+    [NestedCollection] json '$.NestedCollection' AS JSON,
+    [OptionalNested] json '$.OptionalNested' AS JSON,
+    [RequiredNested] json '$.RequiredNested' AS JSON
+) AS [r0]
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
 SELECT [r0].[Id], [r0].[Int], [r0].[Name], [r0].[String], [r0].[NestedCollection], [r0].[OptionalNested], [r0].[RequiredNested]
 FROM [RootEntity] AS [r]
 CROSS APPLY OPENJSON([r].[RelatedCollection], '$') WITH (
@@ -247,6 +266,7 @@ [OptionalNested] nvarchar(max) '$.OptionalNested' AS JSON,
     [RequiredNested] nvarchar(max) '$.RequiredNested' AS JSON
 ) AS [r0]
 """);
+        }
     }
 
     public override async Task SelectMany_nested_collection_on_required_related(QueryTrackingBehavior queryTrackingBehavior)

test/EFCore.SqlServer.FunctionalTests/Query/Associations/ComplexJson/ComplexJsonSetOperationsSqlServerTest.cs

@@ -39,8 +39,38 @@ public override async Task On_related_Select_nested_with_aggregates(QueryTrackin
     {
         await base.On_related_Select_nested_with_aggregates(queryTrackingBehavior);
 
-        AssertSql(
-            """
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+SELECT (
+    SELECT COALESCE(SUM([s].[value]), 0)
+    FROM (
+        SELECT [r0].[NestedCollection] AS [NestedCollection]
+        FROM OPENJSON([r].[RelatedCollection], '$') WITH (
+            [Int] int '$.Int',
+            [NestedCollection] json '$.NestedCollection' AS JSON
+        ) AS [r0]
+        WHERE [r0].[Int] = 8
+        UNION ALL
+        SELECT [r1].[NestedCollection] AS [NestedCollection]
+        FROM OPENJSON([r].[RelatedCollection], '$') WITH (
+            [String] nvarchar(max) '$.String',
+            [NestedCollection] json '$.NestedCollection' AS JSON
+        ) AS [r1]
+        WHERE [r1].[String] = N'foo'
+    ) AS [u]
+    OUTER APPLY (
+        SELECT COALESCE(SUM([n].[Int]), 0) AS [value]
+        FROM OPENJSON([u].[NestedCollection], '$') WITH ([Int] int '$.Int') AS [n]
+    ) AS [s])
+FROM [RootEntity] AS [r]
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
 SELECT (
     SELECT COALESCE(SUM([s].[value]), 0)
     FROM (
@@ -64,6 +94,7 @@ FROM OPENJSON([u].[NestedCollection], '$') WITH ([Int] int '$.Int') AS [n]
     ) AS [s])
 FROM [RootEntity] AS [r]
 """);
+        }
     }
 
     public override async Task On_nested()

test/EFCore.SqlServer.FunctionalTests/Query/Associations/ComplexJson/ComplexJsonStructuralEqualitySqlServerTest.cs

@@ -341,8 +341,39 @@ public override async Task Contains_with_nested_and_composed_operators()
     {
         await base.Contains_with_nested_and_composed_operators();
 
-        AssertSql(
-            """
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+@get_Item_Id='?' (DbType = Int32)
+@entity_equality_get_Item_Id='?' (DbType = Int32)
+@entity_equality_get_Item_Int='?' (DbType = Int32)
+@entity_equality_get_Item_Name='?' (Size = 4000)
+@entity_equality_get_Item_String='?' (Size = 4000)
+@entity_equality_get_Item_NestedCollection='?' (Size = 195)
+@entity_equality_get_Item_OptionalNested='?' (Size = 89)
+@entity_equality_get_Item_RequiredNested='?' (Size = 89)
+
+SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
+FROM [RootEntity] AS [r]
+WHERE EXISTS (
+    SELECT 1
+    FROM OPENJSON([r].[RelatedCollection], '$') WITH (
+        [Id] int '$.Id',
+        [Int] int '$.Int',
+        [Name] nvarchar(max) '$.Name',
+        [String] nvarchar(max) '$.String',
+        [NestedCollection] json '$.NestedCollection' AS JSON,
+        [OptionalNested] json '$.OptionalNested' AS JSON,
+        [RequiredNested] json '$.RequiredNested' AS JSON
+    ) AS [r0]
+    WHERE [r0].[Id] > @get_Item_Id AND [r0].[Id] = @entity_equality_get_Item_Id AND [r0].[Int] = @entity_equality_get_Item_Int AND [r0].[Name] = @entity_equality_get_Item_Name AND [r0].[String] = @entity_equality_get_Item_String AND CAST([r0].[NestedCollection] AS nvarchar(max)) = CAST(@entity_equality_get_Item_NestedCollection AS nvarchar(max)) AND CAST([r0].[OptionalNested] AS nvarchar(max)) = CAST(@entity_equality_get_Item_OptionalNested AS nvarchar(max)) AND CAST([r0].[RequiredNested] AS nvarchar(max)) = CAST(@entity_equality_get_Item_RequiredNested AS nvarchar(max)))
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
 @get_Item_Id='?' (DbType = Int32)
 @entity_equality_get_Item_Id='?' (DbType = Int32)
 @entity_equality_get_Item_Int='?' (DbType = Int32)
@@ -367,6 +398,7 @@ [RequiredNested] nvarchar(max) '$.RequiredNested' AS JSON
     ) AS [r0]
     WHERE [r0].[Id] > @get_Item_Id AND [r0].[Id] = @entity_equality_get_Item_Id AND [r0].[Int] = @entity_equality_get_Item_Int AND [r0].[Name] = @entity_equality_get_Item_Name AND [r0].[String] = @entity_equality_get_Item_String AND [r0].[NestedCollection] = @entity_equality_get_Item_NestedCollection AND [r0].[OptionalNested] = @entity_equality_get_Item_OptionalNested AND [r0].[RequiredNested] = @entity_equality_get_Item_RequiredNested)
 """);
+        }
     }
 
     #endregion Contains