From 00b08d73f9bf75d77375330e06a5a92180005d2f Mon Sep 17 00:00:00 2001
From: Maurycy Markowski <maumar@microsoft.com>
Date: Tue, 24 Sep 2024 09:55:22 -0700
Subject: [PATCH] Fix to #34728 - Split query with
 AsNoTrackingWithIdentityResolution() throws ArgumentOutOfRangeException
 (#34742)

This is a regression introduced in 9.0 when trying to address a different regression (#33073)

Error is caused by a bug in JsonCorrectOrderOfEntitiesForChangeTrackerValidator, specifically it uses the initial SelectExpression to analyze structure of various shaper expressions in the query. Problem is that RelationalSplitCollectionShaperExpression has its own SelectExpression that described the collection - we should use that select expression rather than the parent.
Fix is to update SelectExpression used to process the expression when are processing RelationalSplitCollectionShaperExpression

Fixes #34728
---
 ...sitor.ShaperProcessingExpressionVisitor.cs |   9 +-
 .../Query/AdHocQuerySplittingQueryTestBase.cs | 100 ++++++++++++++++++
 2 files changed, 106 insertions(+), 3 deletions(-)

diff --git a/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs b/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
index f10e922ba68..8b5ff6bd2e2 100644
--- a/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
+++ b/src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs
@@ -3049,7 +3049,7 @@ private sealed class JsonCorrectOrderOfEntitiesForChangeTrackerValidator(SelectE
         {
             private bool _insideCollection;
             private bool _insideInclude;
-
+            private SelectExpression _selectExpression = selectExpression;
             private readonly
                 List<(IEntityType JsonEntityType, List<(IProperty? KeyProperty, int? ConstantKeyValue, int? KeyProjectionIndex)>
                     KeyAccessInfo)> _projectedKeyAccessInfos = [];
@@ -3206,8 +3206,11 @@ protected override Expression VisitExtension(Expression extensionExpression)
                 {
                     var insideCollection = _insideCollection;
                     _insideCollection = true;
+                    var oldSelectExpression = _selectExpression;
+                    _selectExpression = splitCollectionShaperExpression.SelectExpression;
                     Visit(splitCollectionShaperExpression.InnerShaper);
                     _insideCollection = insideCollection;
+                    _selectExpression = oldSelectExpression;
 
                     return splitCollectionShaperExpression;
                 }
@@ -3225,7 +3228,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                         ValueBufferExpression: ProjectionBindingExpression entityProjectionBindingExpression
                     } entityShaperExpression)
                 {
-                    var entityProjection = selectExpression.GetProjection(entityProjectionBindingExpression).GetConstantValue<object>();
+                    var entityProjection = _selectExpression.GetProjection(entityProjectionBindingExpression).GetConstantValue<object>();
 
                     switch (entityProjection)
                     {
@@ -3264,7 +3267,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                     } collectionResultExpression)
                 {
                     var collectionProjection =
-                        selectExpression.GetProjection(collectionProjectionBindingExpression).GetConstantValue<object>();
+                        _selectExpression.GetProjection(collectionProjectionBindingExpression).GetConstantValue<object>();
 
                     switch (collectionProjection)
                     {
diff --git a/test/EFCore.Relational.Specification.Tests/Query/AdHocQuerySplittingQueryTestBase.cs b/test/EFCore.Relational.Specification.Tests/Query/AdHocQuerySplittingQueryTestBase.cs
index 2b3471ebcb4..94e3b92137e 100644
--- a/test/EFCore.Relational.Specification.Tests/Query/AdHocQuerySplittingQueryTestBase.cs
+++ b/test/EFCore.Relational.Specification.Tests/Query/AdHocQuerySplittingQueryTestBase.cs
@@ -360,4 +360,104 @@ public Test(int value)
     }
 
     #endregion
+
+    #region 34728
+
+    [ConditionalTheory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public virtual async Task NoTrackingWithIdentityResolution_split_query_basic(bool async)
+    {
+        var contextFactory = await InitializeAsync<Context34728>(
+            onConfiguring: o => SetQuerySplittingBehavior(o, QuerySplittingBehavior.SplitQuery));
+
+        using var context = contextFactory.CreateContext();
+        var query = context.Set<Context34728.Blog>()
+            .AsNoTrackingWithIdentityResolution()
+            .Select(
+                blog => new
+                {
+                    blog.Id,
+                    Posts = blog.Posts.Select(
+                        blogPost => new 
+                        {
+                            blogPost.Id,
+                            blogPost.Author
+                        }).ToList()
+                });
+
+        var test = async
+            ? await query.ToListAsync()
+            : query.ToList();
+    }
+
+    [ConditionalTheory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public virtual async Task NoTrackingWithIdentityResolution_split_query_complex(bool async)
+    {
+        var contextFactory = await InitializeAsync<Context34728>(
+            onConfiguring: o => SetQuerySplittingBehavior(o, QuerySplittingBehavior.SplitQuery));
+
+        using var context = contextFactory.CreateContext();
+        var query = context.Set<Context34728.Blog>()
+            .AsNoTrackingWithIdentityResolution()
+            .Select(
+                blog => new
+                {
+                    blog.Id,
+                    Posts = blog.Posts.Select(
+                        blogPost => new
+                        {
+                            blogPost.Id,
+                            blogPost.Author
+                        }).ToList(),
+                    Posts2 = blog.Posts.Select(x => new
+                    {
+                        x.Id,
+                        Tags = x.Tags.Select(xx => new
+                        {
+                            xx.Id,
+                            xx.Name,
+                            xx.Name.Length
+                        }).ToList()
+                    }).ToList()
+                });
+
+        var test = async
+            ? await query.ToListAsync()
+            : query.ToList();
+    }
+
+    protected class Context34728(DbContextOptions options) : DbContext(options)
+    {
+        public DbSet<Blog> Tests { get; set; }
+
+        public sealed class Blog
+        {
+            public long Id { get; set; }
+            public string Name { get; set; }
+            public ISet<BlogPost> Posts { get; set; } = new HashSet<BlogPost>();
+        }
+
+        public sealed class BlogPost
+        {
+            public long Id { get; set; }
+            public WebAccount Author { get; set; }
+            public List<Tag> Tags { get; set; }
+        }
+
+        public sealed class WebAccount
+        {
+            public long Id { get; set; }
+        }
+
+        public sealed class Tag
+        {
+            public int Id { get; set; }
+            public string Name { get; set; }
+        }
+    }
+
+    #endregion
 }