WIP: use PK as filters for array aggregates

src/PostgREST/ApiRequest/Types.hs

@@ -144,7 +144,7 @@ type Cast = Text
 type Alias = Text
 type Hint = Text
 
-data AggregateFunction = Sum | Avg | Max | Min | Count | ArrayAgg
+data AggregateFunction = Sum | Avg | Max | Min | Count | ArrayAgg { aaFilters :: [FieldName] }
   deriving (Show, Eq)
 
 data EmbedParam

src/PostgREST/Plan.hs

@@ -332,6 +332,7 @@ readPlan qi@QualifiedIdentifier{..} AppConfig{configDbMaxRows, configDbAggregate
   in
     mapLeft ApiRequestError $
     treeRestrictRange configDbMaxRows (iAction apiRequest) =<<
+    addFiltersToArrayAgg ctx =<<
     hoistSpreadAggFunctions =<<
     validateAggFunctions configDbAggregates =<<
     addRelSelects =<<
@@ -635,7 +636,7 @@ addArrayAggToManySpread (Node rp@ReadPlan{select} forest) =
       shouldAddArrayAgg = spreadRelIsNestedInToMany rp
       fieldToArrayAgg field
         | isJust $ csAggFunction field = field
-        | otherwise = field { csAggFunction = Just ArrayAgg, csAlias = newAlias (csAlias field) (cfName $ csField field) }
+        | otherwise = field { csAggFunction = Just $ ArrayAgg [], csAlias = newAlias (csAlias field) (cfName $ csField field) }
       newAlias alias fieldName = maybe (Just fieldName) pure alias
 
 addRelSelects :: ReadPlanTree -> Either ApiRequestError ReadPlanTree
@@ -676,11 +677,61 @@ generateSpreadSelectFields rp@ReadPlan{select, relSelect} =
     relSelectToSpread :: RelSelectField -> [SpreadSelectField]
     relSelectToSpread (JsonEmbed{rsSelName}) =
       -- The regular embeds that are nested inside spread to-many relationships are also aggregated in an array
-      let (aggFun, alias) = if spreadRelIsNestedInToMany rp then (Just ArrayAgg, Just rsSelName) else (Nothing, Nothing) in
+      let (aggFun, alias) = if spreadRelIsNestedInToMany rp then (Just $ ArrayAgg [], Just rsSelName) else (Nothing, Nothing) in
       [SpreadSelectField { ssSelName = rsSelName, ssSelAggFunction = aggFun, ssSelAggCast = Nothing, ssSelAlias = alias }]
     relSelectToSpread (Spread{rsSpreadSel}) =
       rsSpreadSel
 
+addFiltersToArrayAgg :: ResolverContext -> ReadPlanTree -> Either ApiRequestError ReadPlanTree
+addFiltersToArrayAgg ctx rpt = Right $ applyAddArrayAggFilters ctx [] rpt
+
+-- addFiltersToArrayAgg :: ResolverContext -> ReadPlanTree -> Either ApiRequestError ReadPlanTree
+-- addFiltersToArrayAgg ctx rp@ReadPlan{relSelect} =
+applyAddArrayAggFilters :: ResolverContext -> [(Alias, [CoercibleSelectField])] -> ReadPlanTree -> ReadPlanTree
+applyAddArrayAggFilters ctx pkSelectFields (Node rp@ReadPlan{select, relSelect, relAggAlias} forest) =
+  let newForest = applyAddArrayAggFilters ctx getPKSelectFields <$> forest
+      newSelects
+        | null pkSelectFields = select
+        | otherwise = select ++ fromMaybe mempty (lookup relAggAlias pkSelectFields)
+      newRelSelects
+        | null getPKAliases = relSelect
+        | otherwise = buildPKRelSelect <$> relSelect
+  in Node rp { select = newSelects, relSelect = newRelSelects } newForest
+  where
+    -- Verify if the current node has an array aggregate in the relSelect
+    spreadHasArrAgg Spread{rsSpreadSel} = any (\case Just (ArrayAgg _) -> True; _ -> False; . ssSelAggFunction) rsSpreadSel
+    spreadHasArrAgg _ = False
+    aggSpreads = mapMaybe (\r -> if spreadHasArrAgg r then Just (rsAggAlias r) else Nothing) relSelect
+
+    -- If it has array aggregates, navigate the children nodes to get the PK that will be used as filters for said aggregates
+    allPKSelectFieldsAndAliases = mapMaybe findPKField forest
+    findPKField :: ReadPlanTree -> Maybe ((Alias, [Alias]), (Alias, [CoercibleSelectField]))
+    findPKField (Node ReadPlan{relAggAlias=childAggAlias, from=childTbl} _) =
+      if childAggAlias `elem` aggSpreads
+        then Just $ ((childAggAlias, fst pkFlds), (childAggAlias, snd pkFlds))
+        else Nothing
+      where
+        pkAlias field = childAggAlias <> "_" <> field <> "_pk"
+        toSelectField fld = CoercibleSelectField (resolveOutputField ctx{qi=childTbl} (fld, mempty)) Nothing Nothing Nothing (Just $ pkAlias fld)
+        -- map (\(JoinCondition (_, fld) _) -> (fld, pkAlias fld)) relJoinConds
+        pkFlds = unzip $
+          maybe
+            mempty
+            (map (\pk -> (pkAlias pk, toSelectField pk)) . tablePKCols)
+            (HM.lookup childTbl (tables ctx))
+
+    (getPKAliases, getPKSelectFields) = unzip allPKSelectFieldsAndAliases
+
+    -- Add the PKFields to every ArrayAgg of the respective Spread relSelect
+    buildPKRelSelect rs@Spread{rsAggAlias=rsAlias, rsSpreadSel=rsSel} =
+      case lookup rsAlias getPKAliases of
+        Just pkAliases -> rs{rsSpreadSel= addFilterToArrAgg pkAliases <$> rsSel}
+        _ -> rs
+    buildPKRelSelect rs = rs
+    addFilterToArrAgg pkAliases sel = case ssSelAggFunction sel of
+      Just (ArrayAgg _) -> sel{ssSelAggFunction = Just $ ArrayAgg pkAliases}
+      _                 -> sel
+
 -- When aggregates are present in a ReadPlan that will be spread, we "hoist"
 -- to the highest level possible so that their semantics make sense. For instance,
 -- imagine the user performs the following request:
@@ -763,7 +814,7 @@ hoistIntoRelSelectFields _ r = r
 
 validateAggFunctions :: Bool -> ReadPlanTree -> Either ApiRequestError ReadPlanTree
 validateAggFunctions aggFunctionsAllowed (Node rp@ReadPlan {select} forest)
-  | not aggFunctionsAllowed && any (maybe False (/= ArrayAgg) . csAggFunction) select = Left AggregatesNotAllowed
+  | not aggFunctionsAllowed && any (maybe False (\case ArrayAgg _ -> False; _ -> True) . csAggFunction) select = Left AggregatesNotAllowed
   | otherwise = Node rp <$> traverse (validateAggFunctions aggFunctionsAllowed) forest
 
 addFilters :: ResolverContext -> ApiRequest -> ReadPlanTree -> Either ApiRequestError ReadPlanTree

src/PostgREST/Query/SqlFragment.hs

@@ -295,8 +295,8 @@ pgFmtApplyAggregate (Just agg) aggCast handler snippet =
       -- See https://github.com/PostgREST/postgrest/pull/3640#issuecomment-2334996466
       | otherwise          = "NULLIF(array_agg(" <> snippet <> "),'{null}')"
     aggregatedSnippet = case agg of
-     ArrayAgg -> "COALESCE(" <> fmtArrayAggFunction <> ",'{}')"
-     a        -> convertAggFunction a <> "(" <> snippet <> ")"
+     ArrayAgg _ -> "COALESCE(" <> fmtArrayAggFunction <> ",'{}')"
+     a          -> convertAggFunction a <> "(" <> snippet <> ")"
 
 pgFmtApplyCast :: Maybe Cast -> SQL.Snippet -> SQL.Snippet
 pgFmtApplyCast Nothing snippet = snippet

test/spec/Feature/Query/SpreadQueriesSpec.hs

@@ -135,6 +135,15 @@ spec =
           { matchStatus  = 200
           , matchHeaders = [matchContentTypeJson]
           }
+      it "should return a single null element array, not an empty one, when the row exists but the value happens to be null" $
+        -- See https://github.com/PostgREST/postgrest/pull/3640#discussion_r1786835155
+        get "/managers?select=name,...organizations(organizations:name,referees:referee)&id=eq.1" `shouldRespondWith`
+          [json|[
+            {"name":"Referee Manager","organizations":["Referee Org"],"referees":[null]}
+          ]|]
+          { matchStatus  = 200
+          , matchHeaders = [matchContentTypeJson]
+          }
       it "should work when selecting all columns, aggregating each one of them" $
         get "/factories?select=factory:name,...processes(*)&id=lte.2&order=name" `shouldRespondWith`
           [json|[