Fix type mapping management for JsonScalarExpression (#34663)

Fixes #33582
dotnet · Sep 12, 2024 · 3df883b · 3df883b
1 parent 40a655c
commit 3df883b
Show file tree

Hide file tree

Showing 9 changed files with 117 additions and 13 deletions.
diff --git a/src/EFCore.Relational/Query/RelationalSqlTranslatingExpressionVisitor.cs b/src/EFCore.Relational/Query/RelationalSqlTranslatingExpressionVisitor.cs
@@ -484,7 +484,7 @@ protected override Expression VisitConditional(ConditionalExpression conditional
             || TranslationFailed(conditionalExpression.IfTrue, ifTrue, out var sqlIfTrue)
             || TranslationFailed(conditionalExpression.IfFalse, ifFalse, out var sqlIfFalse)
                 ? QueryCompilationContext.NotTranslatedExpression
-                : _sqlExpressionFactory.Case(new[] { new CaseWhenClause(sqlTest!, sqlIfTrue!) }, sqlIfFalse);
+                : _sqlExpressionFactory.Case([new CaseWhenClause(sqlTest!, sqlIfTrue!)], sqlIfFalse);
     }
 
     /// <inheritdoc />

diff --git a/src/EFCore.Relational/Query/SqlExpressionFactory.cs b/src/EFCore.Relational/Query/SqlExpressionFactory.cs
@@ -357,7 +357,7 @@ private InExpression ApplyTypeMappingOnIn(InExpression inExpression)
 
     private SqlExpression ApplyTypeMappingOnJsonScalar(
         JsonScalarExpression jsonScalarExpression,
-        RelationalTypeMapping? typeMapping)
+        RelationalTypeMapping? elementMapping)
     {
         if (jsonScalarExpression is not { Json: var array, Path: [{ ArrayIndex: { } index }] })
         {
@@ -369,24 +369,28 @@ private SqlExpression ApplyTypeMappingOnJsonScalar(
         var newPath = indexWithTypeMapping == index ? jsonScalarExpression.Path : [new PathSegment(indexWithTypeMapping)];
 
         // If a type mapping is being applied from the outside, it applies to the element resulting from the array indexing operation;
-        // we can infer the array's type mapping from it. Otherwise there's nothing to do but apply the default type mapping to the array.
-        if (typeMapping is null)
+        // we can infer the array's type mapping from it.
+        if (elementMapping is null)
         {
             return new JsonScalarExpression(
-                ApplyDefaultTypeMapping(array),
+                array,
                 newPath,
                 jsonScalarExpression.Type,
-                _typeMappingSource.FindMapping(jsonScalarExpression.Type, Dependencies.Model),
+                jsonScalarExpression.TypeMapping,
                 jsonScalarExpression.IsNullable);
         }
 
-        // TODO: blocked on #30730: we need to be able to construct a JSON collection type mapping based on the element's.
-        // For now, hacking to apply the default type mapping instead.
+        // Resolve the array type mapping for the given element mapping.
+        if (_typeMappingSource.FindMapping(array.Type, Dependencies.Model, elementMapping) is not RelationalTypeMapping arrayMapping)
+        {
+            throw new UnreachableException($"Couldn't find collection type mapping for element type mapping {elementMapping.ClrType.Name}");
+        }
+
         return new JsonScalarExpression(
-            ApplyDefaultTypeMapping(array), // Hack, until #30730
+            ApplyTypeMapping(array, arrayMapping),
             newPath,
             jsonScalarExpression.Type,
-            typeMapping,
+            elementMapping,
             jsonScalarExpression.IsNullable);
     }
 

diff --git a/test/EFCore.Cosmos.FunctionalTests/Query/PrimitiveCollectionsQueryCosmosTest.cs b/test/EFCore.Cosmos.FunctionalTests/Query/PrimitiveCollectionsQueryCosmosTest.cs
@@ -1537,6 +1537,27 @@ FROM root c
 """);
             });
 
+    public override async Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        // Always throws for sync.
+        if (async)
+        {
+            // Member indexer (c.Array[c.SomeMember]) isn't supported by Cosmos
+            var exception = await Assert.ThrowsAsync<CosmosException>(
+                () => base.Parameter_collection_with_type_inference_for_JsonScalarExpression(async));
+
+            Assert.Equal(HttpStatusCode.BadRequest, exception.StatusCode);
+
+            AssertSql(
+                """
+@__values_0='["one","two"]'
+
+SELECT VALUE ((c["Id"] != 0) ? @__values_0[(c["Int"] % 2)] : "foo")
+FROM root c
+""");
+        }
+    }
+
     public override Task Column_collection_Union_parameter_collection(bool async)
         => CosmosTestHelpers.Instance.NoSyncTest(
             async, async a =>

diff --git a/test/EFCore.Specification.Tests/Query/PrimitiveCollectionsQueryTestBase.cs b/test/EFCore.Specification.Tests/Query/PrimitiveCollectionsQueryTestBase.cs
@@ -925,13 +925,24 @@ public virtual Task Inline_collection_Join_ordered_column_collection(bool async)
     [MemberData(nameof(IsAsyncData))]
     public virtual Task Parameter_collection_Concat_column_collection(bool async)
     {
-        var ints = new[] { 11, 111 };
+        int[] ints = [11, 111];
 
         return AssertQuery(
             async,
             ss => ss.Set<PrimitiveCollectionsEntity>().Where(c => ints.Concat(c.Ints).Count() == 2));
     }
 
+    [ConditionalTheory] // #33582
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        string[] values = ["one", "two"];
+
+        return AssertQuery(
+            async,
+            ss => ss.Set<PrimitiveCollectionsEntity>().Select(c => c.Id != 0 ? values[c.Int % 2] : "foo"));
+    }
+
     [ConditionalTheory]
     [MemberData(nameof(IsAsyncData))]
     public virtual Task Column_collection_Union_parameter_collection(bool async)
@@ -1557,8 +1568,8 @@ private static IReadOnlyList<PrimitiveCollectionsEntity> CreatePrimitiveArrayEnt
                     DateTimes = [],
                     Bools = [],
                     Enums = [],
-                    NullableInts = Array.Empty<int?>(),
-                    NullableStrings = Array.Empty<string?>()
+                    NullableInts = [],
+                    NullableStrings = []
                 }
             };
     }

diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQueryOldSqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQueryOldSqlServerTest.cs
@@ -905,6 +905,9 @@ public override Task Inline_collection_Join_ordered_column_collection(bool async
     public override Task Parameter_collection_Concat_column_collection(bool async)
         => AssertCompatibilityLevelTooLow(() => base.Parameter_collection_Concat_column_collection(async));
 
+    public override Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+        => AssertCompatibilityLevelTooLow(() => base.Parameter_collection_Concat_column_collection(async));
+
     public override Task Column_collection_Union_parameter_collection(bool async)
         => AssertCompatibilityLevelTooLow(() => base.Column_collection_Union_parameter_collection(async));
 

diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServer160Test.cs b/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServer160Test.cs
@@ -1507,6 +1507,22 @@ FROM OPENJSON([p].[Ints]) AS [i0]
 """);
     }
 
+    public override async Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        await base.Parameter_collection_with_type_inference_for_JsonScalarExpression(async);
+
+        AssertSql(
+            """
+@__values_0='["one","two"]' (Size = 4000)
+
+SELECT CASE
+    WHEN [p].[Id] <> 0 THEN JSON_VALUE(@__values_0, '$[' + CAST([p].[Int] % 2 AS nvarchar(max)) + ']')
+    ELSE N'foo'
+END
+FROM [PrimitiveCollectionsEntity] AS [p]
+""");
+    }
+
     public override async Task Column_collection_Union_parameter_collection(bool async)
     {
         await base.Column_collection_Union_parameter_collection(async);

diff --git a/.../EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerJsonTypeTest.cs b/.../EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerJsonTypeTest.cs
@@ -1481,6 +1481,22 @@ FROM OPENJSON(CAST([p].[Ints] AS nvarchar(max))) AS [i0]
 """);
     }
 
+    public override async Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        await base.Parameter_collection_with_type_inference_for_JsonScalarExpression(async);
+
+        AssertSql(
+            """
+@__values_0='["one","two"]' (Size = 4000)
+
+SELECT CASE
+    WHEN [p].[Id] <> 0 THEN JSON_VALUE(@__values_0, '$[' + CAST([p].[Int] % 2 AS nvarchar(max)) + ']')
+    ELSE N'foo'
+END
+FROM [PrimitiveCollectionsEntity] AS [p]
+""");
+    }
+
     public override async Task Column_collection_Union_parameter_collection(bool async)
     {
         await base.Column_collection_Union_parameter_collection(async);

diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerTest.cs
@@ -1530,6 +1530,23 @@ FROM OPENJSON([p].[Ints]) AS [i0]
 """);
     }
 
+    [SqlServerCondition(SqlServerCondition.SupportsJsonPathExpressions)]
+    public override async Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        await base.Parameter_collection_with_type_inference_for_JsonScalarExpression(async);
+
+        AssertSql(
+            """
+@__values_0='["one","two"]' (Size = 4000)
+
+SELECT CASE
+    WHEN [p].[Id] <> 0 THEN JSON_VALUE(@__values_0, '$[' + CAST([p].[Int] % 2 AS nvarchar(max)) + ']')
+    ELSE N'foo'
+END
+FROM [PrimitiveCollectionsEntity] AS [p]
+""");
+    }
+
     public override async Task Column_collection_Union_parameter_collection(bool async)
     {
         await base.Column_collection_Union_parameter_collection(async);

diff --git a/test/EFCore.Sqlite.FunctionalTests/Query/PrimitiveCollectionsQuerySqliteTest.cs b/test/EFCore.Sqlite.FunctionalTests/Query/PrimitiveCollectionsQuerySqliteTest.cs
@@ -1489,6 +1489,22 @@ FROM json_each("p"."Ints") AS "i0"
 """);
     }
 
+    public override async Task Parameter_collection_with_type_inference_for_JsonScalarExpression(bool async)
+    {
+        await base.Parameter_collection_with_type_inference_for_JsonScalarExpression(async);
+
+        AssertSql(
+            """
+@__values_0='["one","two"]' (Size = 13)
+
+SELECT CASE
+    WHEN "p"."Id" <> 0 THEN @__values_0 ->> ("p"."Int" % 2)
+    ELSE 'foo'
+END
+FROM "PrimitiveCollectionsEntity" AS "p"
+""");
+    }
+
     public override async Task Column_collection_Union_parameter_collection(bool async)
     {
         await base.Column_collection_Union_parameter_collection(async);