Prune SQL Server OPENJSON's WITH clause

Closes dotnet#32668

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

@@ -60,6 +60,8 @@ public SqlServerOpenJsonExpression(
         IReadOnlyList<ColumnInfo>? columnInfos = null)
         : base(alias, "OPENJSON", schema: null, builtIn: true, new[] { jsonExpression })
     {
+        Check.DebugAssert(columnInfos is null || columnInfos.Count > 0, "Empty column infos not supported, pass null instead");
+
         Path = path;
         ColumnInfos = columnInfos;
     }

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,66 @@
+// 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);
+
+                if (ReferencedColumnMap.TryGetValue(openJson, out var referencedAliases))
+                {
+                    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]);
+                            }
+                        }
+                    }
+
+                    // If we pruned everything, replace the empty list with a null to remove the WITH clause
+                    if (newColumnInfos?.Count == 0)
+                    {
+                        newColumnInfos = null;
+                    }
+
+                    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])
 """);
     }