Skip to content

Commit

Permalink
feat: WIP allow spread operators in to-many relationships
Browse files Browse the repository at this point in the history
  • Loading branch information
laurenceisla committed Sep 15, 2024
1 parent 678103b commit b912c0b
Show file tree
Hide file tree
Showing 10 changed files with 219 additions and 47 deletions.
3 changes: 1 addition & 2 deletions src/PostgREST/ApiRequest/Types.hs
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,6 @@ data ApiRequestError
| PutLimitNotAllowedError
| QueryParamError QPError
| RelatedOrderNotToOne Text Text
| SpreadNotToOne Text Text
| UnacceptableFilter Text
| UnacceptableSchema [Text]
| UnsupportedMethod ByteString
Expand Down Expand Up @@ -145,7 +144,7 @@ type Cast = Text
type Alias = Text
type Hint = Text

data AggregateFunction = Sum | Avg | Max | Min | Count
data AggregateFunction = Sum | Avg | Max | Min | Count | JsonAgg
deriving (Show, Eq)

data EmbedParam
Expand Down
10 changes: 1 addition & 9 deletions src/PostgREST/Error.hs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,6 @@ instance PgrstError ApiRequestError where
status PutLimitNotAllowedError = HTTP.status400
status QueryParamError{} = HTTP.status400
status RelatedOrderNotToOne{} = HTTP.status400
status SpreadNotToOne{} = HTTP.status400
status UnacceptableFilter{} = HTTP.status400
status UnacceptableSchema{} = HTTP.status406
status UnsupportedMethod{} = HTTP.status405
Expand Down Expand Up @@ -176,12 +175,6 @@ instance JSON.ToJSON ApiRequestError where
(Just $ JSON.String $ "'" <> origin <> "' and '" <> target <> "' do not form a many-to-one or one-to-one relationship")
Nothing

toJSON (SpreadNotToOne origin target) = toJsonPgrstError
ApiRequestErrorCode19
("A spread operation on '" <> target <> "' is not possible")
(Just $ JSON.String $ "'" <> origin <> "' and '" <> target <> "' do not form a many-to-one or one-to-one relationship")
Nothing

toJSON (UnacceptableFilter target) = toJsonPgrstError
ApiRequestErrorCode20
("Bad operator on the '" <> target <> "' embedded resource")
Expand Down Expand Up @@ -629,7 +622,7 @@ data ErrorCode
| ApiRequestErrorCode16
| ApiRequestErrorCode17
| ApiRequestErrorCode18
| ApiRequestErrorCode19
-- | ApiRequestErrorCode19 -- no longer used (used to be mapped to SpreadNotToOne)
| ApiRequestErrorCode20
| ApiRequestErrorCode21
| ApiRequestErrorCode22
Expand Down Expand Up @@ -678,7 +671,6 @@ buildErrorCode code = case code of
ApiRequestErrorCode16 -> "PGRST116"
ApiRequestErrorCode17 -> "PGRST117"
ApiRequestErrorCode18 -> "PGRST118"
ApiRequestErrorCode19 -> "PGRST119"
ApiRequestErrorCode20 -> "PGRST120"
ApiRequestErrorCode21 -> "PGRST121"
ApiRequestErrorCode22 -> "PGRST122"
Expand Down
54 changes: 35 additions & 19 deletions src/PostgREST/Plan.hs
Original file line number Diff line number Diff line change
Expand Up @@ -336,10 +336,10 @@ readPlan qi@QualifiedIdentifier{..} AppConfig{configDbMaxRows, configDbAggregate
validateAggFunctions configDbAggregates =<<
addRelSelects =<<
addNullEmbedFilters =<<
validateSpreadEmbeds =<<
addRelatedOrders =<<
addAliases =<<
expandStars ctx =<<
addJsonAggToManySpread =<<
addRels qiSchema (iAction apiRequest) dbRelationships Nothing =<<
addLogicTrees ctx apiRequest =<<
addRanges apiRequest =<<
Expand All @@ -352,7 +352,7 @@ initReadRequest ctx@ResolverContext{qi=QualifiedIdentifier{..}} =
foldr (treeEntry rootDepth) $ Node defReadPlan{from=qi ctx, relName=qiName, depth=rootDepth} []
where
rootDepth = 0
defReadPlan = ReadPlan [] (QualifiedIdentifier mempty mempty) Nothing [] [] allRange mempty Nothing [] Nothing mempty Nothing Nothing False [] rootDepth
defReadPlan = ReadPlan [] (QualifiedIdentifier mempty mempty) Nothing [] [] allRange mempty Nothing [] Nothing mempty Nothing Nothing False False [] rootDepth
treeEntry :: Depth -> Tree SelectItem -> ReadPlanTree -> ReadPlanTree
treeEntry depth (Node si fldForest) (Node q rForest) =
let nxtDepth = succ depth in
Expand Down Expand Up @@ -474,18 +474,18 @@ treeRestrictRange maxRows _ request = pure $ nodeRestrictRange maxRows <$> reque
addRels :: Schema -> Action -> RelationshipsMap -> Maybe ReadPlanTree -> ReadPlanTree -> Either ApiRequestError ReadPlanTree
addRels schema action allRels parentNode (Node rPlan@ReadPlan{relName,relHint,relAlias,depth} forest) =
case parentNode of
Just (Node ReadPlan{from=parentNodeQi, fromAlias=parentAlias} _) ->
Just (Node pr@ReadPlan{from=parentNodeQi, fromAlias=parentAlias} _) ->
let
newReadPlan = (\r ->
let newAlias = Just (qiName (relForeignTable r) <> "_" <> show depth)
aggAlias = qiName (relTable r) <> "_" <> fromMaybe relName relAlias <> "_" <> show depth in
case r of
Relationship{relCardinality=M2M _} -> -- m2m does internal implicit joins that don't need aliasing
rPlan{from=relForeignTable r, relToParent=Just r, relAggAlias=aggAlias, relJoinConds=getJoinConditions Nothing parentAlias r}
rPlan{from=relForeignTable r, relToParent=Just r, relAggAlias=aggAlias, relJoinConds=getJoinConditions Nothing parentAlias r, relIsInToManySpread=spreadRelIsNestedInToMany pr}
ComputedRelationship{} ->
rPlan{from=relForeignTable r, relToParent=Just r{relTableAlias=maybe (relTable r) (QualifiedIdentifier mempty) parentAlias}, relAggAlias=aggAlias, fromAlias=newAlias}
rPlan{from=relForeignTable r, relToParent=Just r{relTableAlias=maybe (relTable r) (QualifiedIdentifier mempty) parentAlias}, relAggAlias=aggAlias, fromAlias=newAlias, relIsInToManySpread=spreadRelIsNestedInToMany pr}
_ ->
rPlan{from=relForeignTable r, relToParent=Just r, relAggAlias=aggAlias, fromAlias=newAlias, relJoinConds=getJoinConditions newAlias parentAlias r}
rPlan{from=relForeignTable r, relToParent=Just r, relAggAlias=aggAlias, fromAlias=newAlias, relJoinConds=getJoinConditions newAlias parentAlias r, relIsInToManySpread=spreadRelIsNestedInToMany pr}
) <$> rel
origin = if depth == 1 -- Only on depth 1 we check if the root(depth 0) has an alias so the sourceCTEName alias can be found as a relationship
then fromMaybe (qiName parentNodeQi) parentAlias
Expand All @@ -509,6 +509,10 @@ addRels schema action allRels parentNode (Node rPlan@ReadPlan{relName,relHint,re
updateForest :: Maybe ReadPlanTree -> Either ApiRequestError [ReadPlanTree]
updateForest rq = addRels schema action allRels rq `traverse` forest

spreadRelIsNestedInToMany :: ReadPlan -> Bool
spreadRelIsNestedInToMany ReadPlan{relIsSpread, relToParent, relIsInToManySpread} =
relIsSpread && (relIsInToManySpread || Just False == (relIsToOne <$> relToParent))

getJoinConditions :: Maybe Alias -> Maybe Alias -> Relationship -> [JoinCondition]
getJoinConditions _ _ ComputedRelationship{} = []
getJoinConditions tblAlias parentAlias Relationship{relTable=qi,relForeignTable=fQi,relCardinality=card} =
Expand Down Expand Up @@ -616,6 +620,24 @@ findRel schema allRels origin target hint =
)
) $ fromMaybe mempty $ HM.lookup (QualifiedIdentifier schema origin, schema) allRels

-- Add JsonAgg aggregates to selected fields that do not have other aggregates and:
-- * Are selected inside a to-many spread relationship
-- * Are selected inside a to-one spread relationship but are nested inside a to-many spread relationship at any level
addJsonAggToManySpread :: ReadPlanTree -> Either ApiRequestError ReadPlanTree
addJsonAggToManySpread (Node rp@ReadPlan{select} forest) =
let newForest = addJsonAggToManySpread `traverse` forest
newSelects
| shouldAddJsonAgg = fieldToJsonAgg <$> select
| otherwise = select
in Node rp { select = newSelects } <$> newForest
where
shouldAddJsonAgg = spreadRelIsNestedInToMany rp
fieldToJsonAgg field
| isJust $ csAggFunction field = field
| otherwise = field { csAggFunction = Just JsonAgg, csAlias = newAlias (csAlias field) (cfName $ csField field) }
newAlias :: Maybe Alias -> FieldName -> Maybe Alias
newAlias _ "*" = Nothing
newAlias alias fieldName = maybe (Just fieldName) pure alias

addRelSelects :: ReadPlanTree -> Either ApiRequestError ReadPlanTree
addRelSelects node@(Node rp forest)
Expand All @@ -628,11 +650,12 @@ addRelSelects node@(Node rp forest)
generateRelSelectField :: ReadPlanTree -> Maybe RelSelectField
generateRelSelectField (Node rp@ReadPlan{relToParent=Just _, relAggAlias, relIsSpread = True} _) =
Just $ Spread { rsSpreadSel = generateSpreadSelectFields rp, rsAggAlias = relAggAlias }
generateRelSelectField (Node ReadPlan{relToParent=Just rel, select, relName, relAlias, relAggAlias, relIsSpread = False} forest) =
generateRelSelectField (Node ReadPlan{relToParent=Just rel, select, relName, relAlias, relAggAlias, relIsSpread = False, relIsInToManySpread} forest) =
Just $ JsonEmbed { rsEmbedMode, rsSelName, rsAggAlias = relAggAlias, rsEmptyEmbed }
where
rsSelName = fromMaybe relName relAlias
rsEmbedMode = if relIsToOne rel then JsonObject else JsonArray
-- If the JsonEmbed is nested in a to-many spread relationship, it will be aggregated at the top. That's why we treat it as `JsonObject`.
rsEmbedMode = if relIsToOne rel || relIsInToManySpread then JsonObject else JsonArray
rsEmptyEmbed = hasOnlyNullEmbed (null select) forest
hasOnlyNullEmbed = foldr checkIfNullEmbed
checkIfNullEmbed :: ReadPlanTree -> Bool -> Bool
Expand All @@ -641,7 +664,7 @@ generateRelSelectField (Node ReadPlan{relToParent=Just rel, select, relName, rel
generateRelSelectField _ = Nothing

generateSpreadSelectFields :: ReadPlan -> [SpreadSelectField]
generateSpreadSelectFields ReadPlan{select, relSelect} =
generateSpreadSelectFields rp@ReadPlan{select, relSelect} =
-- We combine the select and relSelect fields into a single list of SpreadSelectField.
selectSpread ++ relSelectSpread
where
Expand All @@ -653,7 +676,9 @@ generateSpreadSelectFields ReadPlan{select, relSelect} =
relSelectSpread = concatMap relSelectToSpread relSelect
relSelectToSpread :: RelSelectField -> [SpreadSelectField]
relSelectToSpread (JsonEmbed{rsSelName}) =
[SpreadSelectField { ssSelName = rsSelName, ssSelAggFunction = Nothing, ssSelAggCast = Nothing, ssSelAlias = Nothing }]
-- The regular embeds that are nested inside spread to-many relationships are also aggregated usin `json_agg()`
let (aggFun, alias) = if spreadRelIsNestedInToMany rp then (Just JsonAgg, Just rsSelName) else (Nothing, Nothing) in
[SpreadSelectField { ssSelName = rsSelName, ssSelAggFunction = aggFun, ssSelAggCast = Nothing, ssSelAlias = alias }]
relSelectToSpread (Spread{rsSpreadSel}) =
rsSpreadSel

Expand Down Expand Up @@ -906,15 +931,6 @@ resolveLogicTree ctx (Expr b op lts) = CoercibleExpr b op (map (resolveLogicTree
resolveFilter :: ResolverContext -> Filter -> CoercibleFilter
resolveFilter ctx (Filter fld opExpr) = CoercibleFilter{field=resolveQueryInputField ctx fld, opExpr=opExpr}

-- Validates that spread embeds are only done on to-one relationships
validateSpreadEmbeds :: ReadPlanTree -> Either ApiRequestError ReadPlanTree
validateSpreadEmbeds (Node rp@ReadPlan{relToParent=Nothing} forest) = Node rp <$> validateSpreadEmbeds `traverse` forest
validateSpreadEmbeds (Node rp@ReadPlan{relIsSpread,relToParent=Just rel,relName} forest) = do
validRP <- if relIsSpread && not (relIsToOne rel)
then Left $ SpreadNotToOne (qiName $ relTable rel) relName -- TODO using relTable is not entirely right because ReadPlan might have an alias, need to store the parent alias on ReadPlan
else Right rp
Node validRP <$> validateSpreadEmbeds `traverse` forest

-- Find a Node of the Tree and apply a function to it
updateNode :: (a -> ReadPlanTree -> ReadPlanTree) -> (EmbedPath, a) -> Either ApiRequestError ReadPlanTree -> Either ApiRequestError ReadPlanTree
updateNode f ([], a) rr = f a <$> rr
Expand Down
1 change: 1 addition & 0 deletions src/PostgREST/Plan/ReadPlan.hs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ data ReadPlan = ReadPlan
, relHint :: Maybe Hint
, relJoinType :: Maybe JoinType
, relIsSpread :: Bool
, relIsInToManySpread :: Bool
, relSelect :: [RelSelectField]
, depth :: Depth
-- ^ used for aliasing
Expand Down
7 changes: 4 additions & 3 deletions src/PostgREST/Query/SqlFragment.hs
Original file line number Diff line number Diff line change
Expand Up @@ -284,10 +284,11 @@ pgFmtApplyAggregate Nothing _ snippet = snippet
pgFmtApplyAggregate (Just agg) aggCast snippet =
pgFmtApplyCast aggCast aggregatedSnippet
where
convertAggFunction :: AggregateFunction -> SQL.Snippet
aggregatedSnippet = aggFunction <> "(" <> aggArgument <> ")" <> aggFilter
-- Convert from e.g. Sum (the data type) to SUM
convertAggFunction = SQL.sql . BS.map toUpper . BS.pack . show
aggregatedSnippet = convertAggFunction agg <> "(" <> snippet <> ")"
(aggFunction, aggArgument, aggFilter) = case agg of
JsonAgg -> (SQL.sql "json_agg", snippet, mempty)
a -> (SQL.sql . BS.map toUpper . BS.pack $ show a, snippet, mempty)

pgFmtApplyCast :: Maybe Cast -> SQL.Snippet -> SQL.Snippet
pgFmtApplyCast Nothing snippet = snippet
Expand Down
6 changes: 4 additions & 2 deletions test/spec/Feature/Query/AggregateFunctionsSpec.hs
Original file line number Diff line number Diff line change
Expand Up @@ -214,14 +214,16 @@ allowed =
{"name": "Mary", "process_supervisor": [{"name": "Batch", "sum": 220.00}]},
{"name": "John", "process_supervisor": [{"name": "Batch", "sum": 70.00}, {"name": "Mass", "sum": 200.00}]},
{"name": "Peter", "process_supervisor": [{"name": "Batch", "sum": 180.00}, {"name": "Mass", "sum": 120.00}]},
{"name": "Sarah", "process_supervisor": [{"name": "Batch", "sum": 180.00}]}]|]
{"name": "Sarah", "process_supervisor": [{"name": "Batch", "sum": 180.00}]},
{"name": "Jane", "process_supervisor": []}]|]
{ matchHeaders = [matchContentTypeJson] }
get "/supervisors?select=name,process_supervisor(...processes(...process_costs(cost_sum:cost.sum()),...process_categories(category:name)))" `shouldRespondWith`
[json|[
{"name": "Mary", "process_supervisor": [{"category": "Batch", "cost_sum": 220.00}]},
{"name": "John", "process_supervisor": [{"category": "Batch", "cost_sum": 70.00}, {"category": "Mass", "cost_sum": 200.00}]},
{"name": "Peter", "process_supervisor": [{"category": "Batch", "cost_sum": 180.00}, {"category": "Mass", "cost_sum": 120.00}]},
{"name": "Sarah", "process_supervisor": [{"category": "Batch", "cost_sum": 180.00}]}]|]
{"name": "Sarah", "process_supervisor": [{"category": "Batch", "cost_sum": 180.00}]},
{"name": "Jane", "process_supervisor": []}]|]
{ matchHeaders = [matchContentTypeJson] }

context "supports count() aggregate without specifying a field" $ do
Expand Down
Loading

0 comments on commit b912c0b

Please sign in to comment.