Fix to #12173 - Query: Include annotation is not ignored for group by aggregate causing client eval

maumar · maumar · commit 66d135fb3e2e · 2018-06-12T13:07:35.000-07:00
Problem was that when looking for query source on which the include should be applied we were not short-circuiting on subqueries with aggregate result operators that returns a scalar value (count, sum, etc).

Fix is to jump out of the QuerySourceTracingExpressionVisitor when we encounter such a case.
diff --git a/src/EFCore.Specification.Tests/Query/GearsOfWarQueryTestBase.cs b/src/EFCore.Specification.Tests/Query/GearsOfWarQueryTestBase.cs
@@ -5850,6 +5850,136 @@ public virtual void Time_of_day_datetimeoffset()
                      select m.Timeline.TimeOfDay);
         }
 
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_Average()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth ).GroupBy(g => g.Rank).Select(g => g.Average(gg => gg.SquadId)));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_Sum()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.Sum(gg => gg.SquadId)));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_Count()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.Count()));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_LongCount()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.LongCount()));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_Max()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.Max(gg => gg.SquadId)));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Select_Min()
+        {
+            AssertQueryScalar<Gear>(gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.Min(gg => gg.SquadId)));
+        }
+
+        [ConditionalFact]
+        public virtual void GroupBy_Property_Include_Aggregate_with_anonymous_selector()
+        {
+            AssertQuery<Gear>(
+                gs =>
+                    gs.Include(g => g.CityOfBirth).GroupBy(g => g.Nickname).OrderBy(g => g.Key)
+                        .Select(
+                            g => new
+                            {
+                                g.Key,
+                                c = g.Count()
+                            }),
+                assertOrder: true);
+        }
+
+        [ConditionalFact(Skip = "issue #12340")]
+        public virtual void Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector()
+        {
+            AssertQuery<Gear>(
+                gs => gs
+                    .GroupBy(g => g.CityOfBirth)
+                    .OrderBy(g => g.Key.Name)
+                    .Select(g => g.Key)
+                    .Include(c => c.BornGears).ThenInclude(g => g.Weapons),
+                assertOrder: true,
+                elementAsserter: (e, a) =>
+                {
+                    Assert.Equal(e.Name, a.Name);
+
+                    Assert.Equal(e.BornGears == null, a.BornGears == null);
+                    Assert.Equal(e.BornGears.Count(), a.BornGears.Count());
+                });
+        }
+
+        [ConditionalFact(Skip = "issue #12340")]
+        public virtual void Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector_using_EF_Property()
+        {
+            AssertQuery<Gear>(
+                gs => gs
+                    .GroupBy(g => g.CityOfBirth)
+                    .OrderBy(g => g.Key.Name)
+                    .Select(g => EF.Property<City>(g, "Key"))
+                    .Include(c => c.BornGears).ThenInclude(g => g.Weapons),
+                assertOrder: true,
+                elementAsserter: (e, a) =>
+                {
+                    Assert.Equal(e.Name, a.Name);
+
+                    Assert.Equal(e.BornGears == null, a.BornGears == null);
+                    Assert.Equal(e.BornGears.Count(), a.BornGears.Count());
+                });
+        }
+
+        [ConditionalFact]
+        public virtual void Group_by_with_include_with_entity_in_result_selector()
+        {
+            AssertQuery<Gear>(
+                gs =>
+                    gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).OrderBy(g => g.Key)
+                        .Select(
+                            g => new
+                            {
+                                g.Key,
+                                c = g.Count(),
+                                element = g.OrderBy(gg => gg.Nickname).FirstOrDefault()
+                            }),
+                assertOrder: true,
+                elementAsserter: (e, a) =>
+                {
+                    Assert.Equal(e.Key, a.Key);
+                    Assert.Equal(e.c, a.c);
+                    Assert.Equal(e.element.Nickname, a.element.Nickname);
+                    Assert.Equal(e.element.CityOfBirth == null, a.element.CityOfBirth == null);
+                    if (e.element.CityOfBirth != null)
+                    {
+                        Assert.Equal(e.element.CityOfBirth.Name, a.element.CityOfBirth.Name);
+                    }
+                });
+        }
+
+        [ConditionalFact]
+        public virtual void Include_with_group_by_and_FirstOrDefault_gets_properly_applied()
+        {
+            var expectedIncludes = new List<IExpectedInclude>
+            {
+                new ExpectedInclude<Gear>(e => e.CityOfBirth, "CityOfBirth"),
+                new ExpectedInclude<Officer>(e => e.CityOfBirth, "CityOfBirth"),
+            };
+
+            AssertIncludeQuery<Gear>(
+                gs => gs.Include(g => g.CityOfBirth).GroupBy(g => g.Rank).Select(g => g.FirstOrDefault(gg => gg.HasSoulPatch)),
+                expectedIncludes);
+        }
+
         // Remember to add any new tests to Async version of this test class
 
         protected GearsOfWarContext CreateContext() => Fixture.CreateContext();
diff --git a/src/EFCore/Query/ExpressionVisitors/Internal/QuerySourceTracingExpressionVisitor.cs b/src/EFCore/Query/ExpressionVisitors/Internal/QuerySourceTracingExpressionVisitor.cs
@@ -1,6 +1,8 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Linq.Expressions;
 using JetBrains.Annotations;
@@ -100,6 +102,16 @@ protected override Expression VisitQuerySourceReference(QuerySourceReferenceExpr
             return expression;
         }
 
+        private static readonly ISet<Type> _aggregateResultOperators = new HashSet<Type>
+        {
+            typeof(AverageResultOperator),
+            typeof(CountResultOperator),
+            typeof(LongCountResultOperator),
+            typeof(MaxResultOperator),
+            typeof(MinResultOperator),
+            typeof(SumResultOperator)
+        };
+
         /// <summary>
         ///     This API supports the Entity Framework Core infrastructure and is not intended to be used
         ///     directly from your code. This API may change or be removed in future releases.
@@ -111,6 +123,11 @@ protected override Expression VisitSubQuery(SubQueryExpression subQueryExpressio
                 return subQueryExpression;
             }
 
+            if (subQueryExpression.QueryModel.ResultOperators.Any(ro => _aggregateResultOperators.Contains(ro.GetType())))
+            {
+                return subQueryExpression;
+            }
+
             _originGroupByQueryModel
                 = subQueryExpression.QueryModel.ResultOperators
                     .Any(ro => ro is GroupResultOperator)
diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/GearsOfWarQuerySqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/GearsOfWarQuerySqlServerTest.cs
@@ -7614,6 +7614,124 @@ public override void Time_of_day_datetimeoffset()
 FROM [Missions] AS [m]");
         }
 
+        public override void GroupBy_Property_Include_Select_Average()
+        {
+            base.GroupBy_Property_Include_Select_Average();
+
+            AssertSql(
+                @"SELECT AVG(CAST([g].[SquadId] AS float))
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Select_Sum()
+        {
+            base.GroupBy_Property_Include_Select_Sum();
+
+            AssertSql(
+                @"SELECT SUM([g].[SquadId])
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Select_Count()
+        {
+            base.GroupBy_Property_Include_Select_Count();
+
+            AssertSql(
+                @"SELECT COUNT(*)
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Select_LongCount()
+        {
+            base.GroupBy_Property_Include_Select_LongCount();
+
+            AssertSql(
+                @"SELECT COUNT_BIG(*)
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Select_Min()
+        {
+            base.GroupBy_Property_Include_Select_Min();
+
+            AssertSql(
+                @"SELECT MIN([g].[SquadId])
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Aggregate_with_anonymous_selector()
+        {
+            base.GroupBy_Property_Include_Aggregate_with_anonymous_selector();
+
+            AssertSql(
+                @"SELECT [g].[Nickname] AS [Key], COUNT(*) AS [c]
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Nickname]
+ORDER BY [Key]");
+        }
+
+        public override void Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector()
+        {
+            base.Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector();
+
+            AssertSql(
+                @"");
+        }
+
+        public override void Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector_using_EF_Property()
+        {
+            base.Group_by_entity_key_with_include_on_that_entity_with_key_in_result_selector_using_EF_Property();
+
+            AssertSql(
+                @"");
+        }
+
+        public override void Group_by_with_include_with_entity_in_result_selector()
+        {
+            base.Group_by_with_include_with_entity_in_result_selector();
+
+            AssertSql(
+                @"SELECT [g].[Nickname], [g].[SquadId], [g].[AssignedCityName], [g].[CityOrBirthName], [g].[Discriminator], [g].[FullName], [g].[HasSoulPatch], [g].[LeaderNickname], [g].[LeaderSquadId], [g].[Rank], [g.CityOfBirth].[Name], [g.CityOfBirth].[Location]
+FROM [Gears] AS [g]
+INNER JOIN [Cities] AS [g.CityOfBirth] ON [g].[CityOrBirthName] = [g.CityOfBirth].[Name]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+ORDER BY [g].[Rank]");
+        }
+
+        public override void GroupBy_Property_Include_Select_Max()
+        {
+            base.GroupBy_Property_Include_Select_Max();
+
+            AssertSql(
+                @"SELECT MAX([g].[SquadId])
+FROM [Gears] AS [g]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+GROUP BY [g].[Rank]");
+        }
+
+        public override void Include_with_group_by_and_FirstOrDefault_gets_properly_applied()
+        {
+            base.Include_with_group_by_and_FirstOrDefault_gets_properly_applied();
+
+            AssertSql(
+                @"SELECT [g].[Nickname], [g].[SquadId], [g].[AssignedCityName], [g].[CityOrBirthName], [g].[Discriminator], [g].[FullName], [g].[HasSoulPatch], [g].[LeaderNickname], [g].[LeaderSquadId], [g].[Rank], [g.CityOfBirth].[Name], [g.CityOfBirth].[Location]
+FROM [Gears] AS [g]
+INNER JOIN [Cities] AS [g.CityOfBirth] ON [g].[CityOrBirthName] = [g.CityOfBirth].[Name]
+WHERE [g].[Discriminator] IN (N'Officer', N'Gear')
+ORDER BY [g].[Rank]");
+        }
+
         private void AssertSql(params string[] expected)
             => Fixture.TestSqlLoggerFactory.AssertBaseline(expected);
 
