Prune SQL Server OPENJSON's WITH clause (dotnet#32673)

Closes dotnet#32668

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

@@ -60,6 +60,11 @@ public SqlServerOpenJsonExpression(
         IReadOnlyList<ColumnInfo>? columnInfos = null)
         : base(alias, "OPENJSON", schema: null, builtIn: true, new[] { jsonExpression })
     {
+        if (columnInfos?.Count == 0)
+        {
+            columnInfos = null;
+        }
+
         Path = path;
         ColumnInfos = columnInfos;
     }
@@ -130,12 +135,19 @@ public virtual SqlServerOpenJsonExpression Update(
         SqlExpression jsonExpression,
         IReadOnlyList<PathSegment>? path,
         IReadOnlyList<ColumnInfo>? columnInfos = null)
-        => jsonExpression == JsonExpression
+    {
+        if (columnInfos?.Count == 0)
+        {
+            columnInfos = null;
+        }
+
+        return jsonExpression == JsonExpression
             && (ReferenceEquals(path, Path) || path is not null && Path is not null && path.SequenceEqual(Path))
             && (ReferenceEquals(columnInfos, ColumnInfos)
                 || columnInfos is not null && ColumnInfos is not null && columnInfos.SequenceEqual(ColumnInfos))
                 ? this
                 : new SqlServerOpenJsonExpression(Alias, jsonExpression, path, columnInfos);
+    }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Diagnostics.CodeAnalysis;
+using Microsoft.EntityFrameworkCore.Query.Internal;
 using Microsoft.EntityFrameworkCore.Query.SqlExpressions;
 using Microsoft.EntityFrameworkCore.SqlServer.Internal;
 
@@ -17,6 +18,7 @@ public class SqlServerQueryTranslationPostprocessor : RelationalQueryTranslation
 {
     private readonly SqlServerJsonPostprocessor _jsonPostprocessor;
     private readonly SkipWithoutOrderByInSplitQueryVerifier _skipWithoutOrderByInSplitQueryVerifier = new();
+    private readonly SqlServerSqlTreePruner _pruner = new();
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -50,6 +52,15 @@ public override Expression Process(Expression query)
         return query;
     }
 
+    /// <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
+    ///     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>
+    protected override Expression Prune(Expression query)
+        => _pruner.Prune(query);
+
     private sealed class SkipWithoutOrderByInSplitQueryVerifier : ExpressionVisitor
     {
         [return: NotNullIfNotNull("expression")]

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

@@ -0,0 +1,63 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.EntityFrameworkCore.Query.SqlExpressions;
+using ColumnInfo = Microsoft.EntityFrameworkCore.SqlServer.Query.Internal.SqlServerOpenJsonExpression.ColumnInfo;
+
+namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal;
+
+/// <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
+///     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 class SqlServerSqlTreePruner : SqlTreePruner
+{
+    /// <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
+    ///     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>
+    protected override Expression VisitExtension(Expression node)
+    {
+        switch (node)
+        {
+            case SqlServerOpenJsonExpression { ColumnInfos: IReadOnlyList<ColumnInfo> columnInfos } openJson:
+                var visitedJson = (SqlExpression)Visit(openJson.JsonExpression);
+
+#pragma warning disable EF1001 // ReferencedColumnMap is pubternal; should be made protected
+                if (ReferencedColumnMap.TryGetValue(openJson, out var referencedAliases))
+#pragma warning restore EF1001
+                {
+                    List<ColumnInfo>? newColumnInfos = null;
+
+                    for (var i = 0; i < columnInfos.Count; i++)
+                    {
+                        if (referencedAliases.Contains(columnInfos[i].Name))
+                        {
+                            newColumnInfos?.Add(columnInfos[i]);
+                        }
+                        else if (newColumnInfos is null)
+                        {
+                            newColumnInfos = new();
+                            for (var j = 0; j < i; j++)
+                            {
+                                newColumnInfos.Add(columnInfos[j]);
+                            }
+                        }
+                    }
+
+                    // Not that if we pruned everything, the WITH clause gets removed entirely
+                    return openJson.Update(visitedJson, openJson.Path, newColumnInfos ?? openJson.ColumnInfos);
+                }
+
+                // There are no references to the OPENJSON expression; remove the WITH clause entirely
+                return openJson.Update(visitedJson, openJson.Path);
+
+            default:
+                return base.VisitExtension(node);
+        }
+    }
+}

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

@@ -1011,16 +1011,7 @@ public override async Task Json_collection_Any_with_predicate(bool async)
 FROM [JsonEntitiesBasic] AS [j]
 WHERE EXISTS (
     SELECT 1
-    FROM OPENJSON([j].[OwnedReferenceRoot], '$.OwnedCollectionBranch') WITH (
-        [Date] datetime2 '$.Date',
-        [Enum] int '$.Enum',
-        [Enums] nvarchar(max) '$.Enums' AS JSON,
-        [Fraction] decimal(18,2) '$.Fraction',
-        [NullableEnum] int '$.NullableEnum',
-        [NullableEnums] nvarchar(max) '$.NullableEnums' AS JSON,
-        [OwnedCollectionLeaf] nvarchar(max) '$.OwnedCollectionLeaf' AS JSON,
-        [OwnedReferenceLeaf] nvarchar(max) '$.OwnedReferenceLeaf' AS JSON
-    ) AS [o]
+    FROM OPENJSON([j].[OwnedReferenceRoot], '$.OwnedCollectionBranch') WITH ([OwnedReferenceLeaf] nvarchar(max) '$.OwnedReferenceLeaf' AS JSON) AS [o]
     WHERE JSON_VALUE([o].[OwnedReferenceLeaf], '$.SomethingSomething') = N'e1_r_c1_r')
 """);
     }
@@ -1078,11 +1069,7 @@ SELECT JSON_VALUE([o].[OwnedReferenceLeaf], '$.SomethingSomething') AS [c], [o].
         FROM OPENJSON([j].[OwnedReferenceRoot], '$.OwnedCollectionBranch') WITH (
             [Date] datetime2 '$.Date',
             [Enum] int '$.Enum',
-            [Enums] nvarchar(max) '$.Enums' AS JSON,
             [Fraction] decimal(18,2) '$.Fraction',
-            [NullableEnum] int '$.NullableEnum',
-            [NullableEnums] nvarchar(max) '$.NullableEnums' AS JSON,
-            [OwnedCollectionLeaf] nvarchar(max) '$.OwnedCollectionLeaf' AS JSON,
             [OwnedReferenceLeaf] nvarchar(max) '$.OwnedReferenceLeaf' AS JSON
         ) AS [o]
         ORDER BY [o].[Date] DESC
@@ -1130,26 +1117,10 @@ public override async Task Json_collection_within_collection_Count(bool async)
 FROM [JsonEntitiesBasic] AS [j]
 WHERE EXISTS (
     SELECT 1
-    FROM OPENJSON([j].[OwnedCollectionRoot], '$') WITH (
-        [Name] nvarchar(max) '$.Name',
-        [Names] nvarchar(max) '$.Names' AS JSON,
-        [Number] int '$.Number',
-        [Numbers] nvarchar(max) '$.Numbers' AS JSON,
-        [OwnedCollectionBranch] nvarchar(max) '$.OwnedCollectionBranch' AS JSON,
-        [OwnedReferenceBranch] nvarchar(max) '$.OwnedReferenceBranch' AS JSON
-    ) AS [o]
+    FROM OPENJSON([j].[OwnedCollectionRoot], '$') WITH ([OwnedCollectionBranch] nvarchar(max) '$.OwnedCollectionBranch' AS JSON) AS [o]
     WHERE (
         SELECT COUNT(*)
-        FROM OPENJSON([o].[OwnedCollectionBranch], '$') WITH (
-            [Date] datetime2 '$.Date',
-            [Enum] int '$.Enum',
-            [Enums] nvarchar(max) '$.Enums' AS JSON,
-            [Fraction] decimal(18,2) '$.Fraction',
-            [NullableEnum] int '$.NullableEnum',
-            [NullableEnums] nvarchar(max) '$.NullableEnums' AS JSON,
-            [OwnedCollectionLeaf] nvarchar(max) '$.OwnedCollectionLeaf' AS JSON,
-            [OwnedReferenceLeaf] nvarchar(max) '$.OwnedReferenceLeaf' AS JSON
-        ) AS [o0]) = 2)
+        FROM OPENJSON([o].[OwnedCollectionBranch], '$') AS [o0]) = 2)
 """);
     }
 
@@ -1161,14 +1132,7 @@ public override async Task Json_collection_in_projection_with_composition_count(
             """
 SELECT (
     SELECT COUNT(*)
-    FROM OPENJSON([j].[OwnedCollectionRoot], '$') WITH (
-        [Name] nvarchar(max) '$.Name',
-        [Names] nvarchar(max) '$.Names' AS JSON,
-        [Number] int '$.Number',
-        [Numbers] nvarchar(max) '$.Numbers' AS JSON,
-        [OwnedCollectionBranch] nvarchar(max) '$.OwnedCollectionBranch' AS JSON,
-        [OwnedReferenceBranch] nvarchar(max) '$.OwnedReferenceBranch' AS JSON
-    ) AS [o])
+    FROM OPENJSON([j].[OwnedCollectionRoot], '$') AS [o])
 FROM [JsonEntitiesBasic] AS [j]
 ORDER BY [j].[Id]
 """);

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

@@ -562,7 +562,7 @@ public override async Task Column_collection_Count_method(bool async)
 FROM [PrimitiveCollectionsEntity] AS [p]
 WHERE (
     SELECT COUNT(*)
-    FROM OPENJSON([p].[Ints]) WITH ([value] int '$') AS [i]) = 2
+    FROM OPENJSON([p].[Ints]) AS [i]) = 2
 """);
     }
 
@@ -576,7 +576,7 @@ public override async Task Column_collection_Length(bool async)
 FROM [PrimitiveCollectionsEntity] AS [p]
 WHERE (
     SELECT COUNT(*)
-    FROM OPENJSON([p].[Ints]) WITH ([value] int '$') AS [i]) = 2
+    FROM OPENJSON([p].[Ints]) AS [i]) = 2
 """);
     }
 
@@ -650,7 +650,7 @@ public override async Task Non_nullable_reference_column_collection_index_equals
 FROM [PrimitiveCollectionsEntity] AS [p]
 WHERE EXISTS (
     SELECT 1
-    FROM OPENJSON([p].[Strings]) WITH ([value] nvarchar(max) '$') AS [s]) AND JSON_VALUE([p].[Strings], '$[1]') = [p].[NullableString]
+    FROM OPENJSON([p].[Strings]) AS [s]) AND JSON_VALUE([p].[Strings], '$[1]') = [p].[NullableString]
 """);
     }
 
@@ -790,7 +790,7 @@ public override async Task Column_collection_Any(bool async)
 FROM [PrimitiveCollectionsEntity] AS [p]
 WHERE EXISTS (
     SELECT 1
-    FROM OPENJSON([p].[Ints]) WITH ([value] int '$') AS [i])
+    FROM OPENJSON([p].[Ints]) AS [i])
 """);
     }
 
@@ -881,10 +881,10 @@ FROM [PrimitiveCollectionsEntity] AS [p]
     SELECT COUNT(*)
     FROM (
         SELECT 1 AS empty
-        FROM OPENJSON(@__ints_0) WITH ([value] int '$') AS [i]
+        FROM OPENJSON(@__ints_0) AS [i]
         UNION ALL
         SELECT 1 AS empty
-        FROM OPENJSON([p].[Ints]) WITH ([value] int '$') AS [i0]
+        FROM OPENJSON([p].[Ints]) AS [i0]
     ) AS [t]) = 2
 """);
     }