Fix nullability processing for aggregate PostgresFunctionExpression (#2396)

Author: roji

src/EFCore.PG/Query/Expressions/Internal/PostgresFunctionExpression.cs

@@ -217,7 +217,9 @@ public override SqlFunctionExpression Update(SqlExpression? instance, IReadOnlyL
             : this;
     }
 
-    public virtual SqlFunctionExpression UpdateAggregateComponents(SqlExpression? predicate, IReadOnlyList<OrderingExpression> orderings)
+    public virtual PostgresFunctionExpression UpdateAggregateComponents(
+        SqlExpression? predicate,
+        IReadOnlyList<OrderingExpression> orderings)
     {
         return predicate != AggregatePredicate || orderings != AggregateOrderings
             ? new PostgresFunctionExpression(

src/EFCore.PG/Query/Internal/NpgsqlSqlNullabilityProcessor.cs

@@ -36,8 +36,6 @@ PostgresArrayIndexExpression arrayIndexExpression
                 => VisitArrayIndex(arrayIndexExpression, allowOptimizedExpansion, out nullable),
             PostgresBinaryExpression binaryExpression
                 => VisitBinary(binaryExpression, allowOptimizedExpansion, out nullable),
-            PostgresFunctionExpression postgresFunctionExpression
-                => VisitPostgresFunction(postgresFunctionExpression, allowOptimizedExpansion, out nullable),
             PostgresILikeExpression ilikeExpression
                 => VisitILike(ilikeExpression, allowOptimizedExpansion, out nullable),
             PostgresJsonTraversalExpression postgresJsonTraversalExpression
@@ -51,6 +49,8 @@ PostgresRowValueExpression postgresRowValueExpression
             PostgresUnknownBinaryExpression postgresUnknownBinaryExpression
                 => VisitUnknownBinary(postgresUnknownBinaryExpression, allowOptimizedExpansion, out nullable),
 
+            // PostgresFunctionExpression is visited via the SqlFunctionExpression override below
+
             _ => base.VisitCustomSqlExpression(sqlExpression, allowOptimizedExpansion, out nullable)
         };
 
@@ -182,46 +182,68 @@ protected virtual SqlExpression VisitBinary(
         return binaryExpression.Update(left, right);
     }
 
-    protected virtual SqlExpression VisitPostgresFunction(
-        PostgresFunctionExpression functionExpression,
+    protected override SqlExpression VisitSqlFunction(
+        SqlFunctionExpression sqlFunctionExpression,
         bool allowOptimizedExpansion,
         out bool nullable)
     {
         // PostgresFunctionExpression extends SqlFunctionExpression, and adds aggregate predicate and ordering expressions to that.
         // First call the base VisitSqlFunction to visit the arguments
-        var visitedBase =
-            (PostgresFunctionExpression)base.VisitSqlFunction(functionExpression, allowOptimizedExpansion, out nullable);
+        var visitedBase = base.VisitSqlFunction(sqlFunctionExpression, allowOptimizedExpansion, out nullable);
+
+        // base.VisitSqlFunction has some special logic for SUM which wraps it in a COALESCE
+        // (see https://github.com/dotnet/efcore/issues/28158), so we need some special handling to properly visit the
+        // PostgresFunctionExpression it wraps.
+        if (sqlFunctionExpression.IsBuiltIn
+            && string.Equals(sqlFunctionExpression.Name, "SUM", StringComparison.OrdinalIgnoreCase)
+            && visitedBase is SqlFunctionExpression { Name: "COALESCE", Arguments: { } } coalesceExpression
+            && coalesceExpression.Arguments[0] is PostgresFunctionExpression wrappedFunctionExpression)
+        {
+            var visitedArguments = coalesceExpression.Arguments!.ToArray();
+            visitedArguments[0] = VisitPostgresFunctionComponents(wrappedFunctionExpression);
 
-        var aggregateChanged = false;
+            return coalesceExpression.Update(coalesceExpression.Instance, visitedArguments);
+        }
 
-        var visitedAggregatePredicate = Visit(functionExpression.AggregatePredicate, allowOptimizedExpansion: true, out _);
-        aggregateChanged |= visitedAggregatePredicate != functionExpression.AggregatePredicate;
+        return visitedBase is PostgresFunctionExpression pgFunctionExpression
+            ? VisitPostgresFunctionComponents(pgFunctionExpression)
+            : visitedBase;
 
-        OrderingExpression[]? visitedOrderings = null;
-        for (var i = 0; i < functionExpression.AggregateOrderings.Count; i++)
+        PostgresFunctionExpression VisitPostgresFunctionComponents(PostgresFunctionExpression pgFunctionExpression)
         {
-            var ordering = functionExpression.AggregateOrderings[i];
-            var visitedOrdering = ordering.Update(Visit(ordering.Expression, out _));
-            if (visitedOrdering != ordering && visitedOrderings is null)
+            var aggregateChanged = false;
+
+            var visitedAggregatePredicate = Visit(pgFunctionExpression.AggregatePredicate, allowOptimizedExpansion: true, out _);
+            aggregateChanged |= visitedAggregatePredicate != pgFunctionExpression.AggregatePredicate;
+
+            OrderingExpression[]? visitedOrderings = null;
+            for (var i = 0; i < pgFunctionExpression.AggregateOrderings.Count; i++)
             {
-                visitedOrderings = new OrderingExpression[functionExpression.AggregateOrderings.Count];
-                for (var j = 0; j < i; j++)
+                var ordering = pgFunctionExpression.AggregateOrderings[i];
+                var visitedOrdering = ordering.Update(Visit(ordering.Expression, out _));
+                if (visitedOrdering != ordering && visitedOrderings is null)
                 {
-                    visitedOrderings[j] = functionExpression.AggregateOrderings[j];
+                    visitedOrderings = new OrderingExpression[pgFunctionExpression.AggregateOrderings.Count];
+                    for (var j = 0; j < i; j++)
+                    {
+                        visitedOrderings[j] = pgFunctionExpression.AggregateOrderings[j];
+                    }
+
+                    aggregateChanged = true;
                 }
 
-                aggregateChanged = true;
+                if (visitedOrderings is not null)
+                {
+                    visitedOrderings[i] = visitedOrdering;
+                }
             }
 
-            if (visitedOrderings is not null)
-            {
-                visitedOrderings[i] = visitedOrdering;
-            }
+            return aggregateChanged
+                ? pgFunctionExpression.UpdateAggregateComponents(
+                    visitedAggregatePredicate,
+                    visitedOrderings ?? pgFunctionExpression.AggregateOrderings)
+                : pgFunctionExpression;
         }
-
-        return aggregateChanged
-            ? visitedBase.UpdateAggregateComponents(visitedAggregatePredicate, visitedOrderings ?? functionExpression.AggregateOrderings)
-            : visitedBase;
     }
 
     protected virtual SqlExpression VisitILike(