From bca4b0e8d2fc7bfe647071fef518a44c505093fb Mon Sep 17 00:00:00 2001 From: Laurence Isla Date: Sat, 19 Oct 2024 11:23:19 -0500 Subject: [PATCH] WIP: use PK as filters for array aggregates --- src/PostgREST/ApiRequest/Types.hs | 2 +- src/PostgREST/Plan.hs | 57 ++++++++++++++++++-- src/PostgREST/Query/SqlFragment.hs | 4 +- test/spec/Feature/Query/SpreadQueriesSpec.hs | 9 ++++ 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/src/PostgREST/ApiRequest/Types.hs b/src/PostgREST/ApiRequest/Types.hs index 9dbd369da1..e98133016a 100644 --- a/src/PostgREST/ApiRequest/Types.hs +++ b/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 diff --git a/src/PostgREST/Plan.hs b/src/PostgREST/Plan.hs index 397f7cf33c..623fccba69 100644 --- a/src/PostgREST/Plan.hs +++ b/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 diff --git a/src/PostgREST/Query/SqlFragment.hs b/src/PostgREST/Query/SqlFragment.hs index c9d9b2cc26..1019c4819c 100644 --- a/src/PostgREST/Query/SqlFragment.hs +++ b/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 diff --git a/test/spec/Feature/Query/SpreadQueriesSpec.hs b/test/spec/Feature/Query/SpreadQueriesSpec.hs index 27eeef9117..579bce286e 100644 --- a/test/spec/Feature/Query/SpreadQueriesSpec.hs +++ b/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|[