refactor: add the returningCols function

* separate fieldNames from getting fkCols

* Put binaryField inside readSqlParts

* Move scalar proc logic to binaryField

* Move logic for the "SELECT *" default to DbRequestBuilder

src/PostgREST/ApiRequest.hs

@@ -104,7 +104,7 @@ data ApiRequest = ApiRequest {
   -- | &and and &or parameters used for complex boolean logic
   , iLogic                :: [(Text, Text)]
   -- | &select parameter used to shape the response
-  , iSelect               :: Text
+  , iSelect               :: Maybe Text
   -- | &columns parameter used to shape the payload
   , iColumns              :: Maybe Text
   -- | &order parameters for each level
@@ -145,7 +145,7 @@ userApiRequest schema rootSpec req reqBody
                                | otherwise                         -> Nothing
       , iFilters = filters
       , iLogic = [(toS k, toS $ fromJust v) | (k,v) <- qParams, isJust v, endingIn ["and", "or"] k ]
-      , iSelect = toS $ fromMaybe "*" $ join $ lookup "select" qParams
+      , iSelect = toS <$> join (lookup "select" qParams)
       , iColumns = columns
       , iOrder = [(toS k, toS $ fromJust v) | (k,v) <- qParams, isJust v, endingIn ["order"] k ]
       , iCanonicalQS = toS $ urlEncodeVars

src/PostgREST/App.hs

@@ -55,10 +55,10 @@ import PostgREST.Error            (PgError (..), SimpleError (..),
 import PostgREST.Middleware
 import PostgREST.OpenAPI
 import PostgREST.Parsers          (pRequestColumns)
-import PostgREST.QueryBuilder     (limitedQuery,
+import PostgREST.QueryBuilder     (limitedQuery, mutateRequestToQuery,
-                                   requestToCallProcQuery,
                                    readRequestToCountQuery,
-                                   readRequestToQuery, mutateRequestToQuery)
+                                   readRequestToQuery,
+                                   requestToCallProcQuery)
 import PostgREST.RangeQuery       (allRange, contentRangeH,
                                    rangeStatusHeader)
 import PostgREST.Statements       (callProcStatement,
@@ -119,17 +119,16 @@ transactionMode proc action =
 
 app :: DbStructure -> Maybe ProcDescription -> S.Set FieldName -> AppConfig -> ApiRequest -> H.Transaction Response
 app dbStructure proc cols conf apiRequest =
+  let rawContentTypes = (decodeContentType <$> configRawMediaTypes conf) `L.union` [ CTOctetStream, CTTextPlain ] in
   case responseContentTypeOrError (iAccepts apiRequest) rawContentTypes (iAction apiRequest) (iTarget apiRequest) of
     Left errorResponse -> return errorResponse
     Right contentType ->
       case (iAction apiRequest, iTarget apiRequest, iPayload apiRequest) of
 
         (ActionRead headersOnly, TargetIdent (QualifiedIdentifier _ tName), Nothing) ->
-          let partsField = (,) <$> readSqlParts tName
+          case readSqlParts tName of
-                <*> (binaryField contentType rawContentTypes =<< fldNames tName) in
-          case partsField of
             Left errorResponse -> return errorResponse
-            Right ((q, cq), bField) -> do
+            Right (q, cq, bField) -> do
               let cQuery = if estimatedCount
                              then limitedQuery cq ((+ 1) <$> maxRows) -- LIMIT maxRows + 1 so we can determine below that maxRows was surpassed
                              else cq
@@ -283,15 +282,10 @@ app dbStructure proc cols conf apiRequest =
               return $ responseLBS status200 [allOrigins, allowH] mempty
 
         (ActionInvoke invMethod, TargetProc qi@(QualifiedIdentifier _ pName) _, Just pJson) ->
-          let returnsScalar = maybe False procReturnsScalar proc
+          let tName = fromMaybe pName $ procTableName =<< proc in
-              tName = fromMaybe pName $ procTableName =<< proc
+          case readSqlParts tName of
-              rpcBinaryField = if returnsScalar
-                                 then Right Nothing
-                                 else binaryField contentType rawContentTypes =<< fldNames tName
-              parts = (,) <$> readSqlParts tName <*> rpcBinaryField in
-          case parts of
             Left errorResponse -> return errorResponse
-            Right ((q, cq), bField) -> do
+            Right (q, cq, bField) -> do
               let
                 preferParams = iPreferParameters apiRequest
                 pq = requestToCallProcQuery qi (specifiedProcArgs cols proc) returnsScalar preferParams
@@ -337,19 +331,26 @@ app dbStructure proc cols conf apiRequest =
         plannedCount = iPreferCount apiRequest == Just PlannedCount
         shouldCount = exactCount || estimatedCount
         topLevelRange = iTopLevelRange apiRequest
-      readReq tableName = readRequest schema tableName maxRows (dbRelations dbStructure) apiRequest
+        returnsScalar = maybe False procReturnsScalar proc
-      fldNames tableName = fieldNames <$> readReq tableName
+
-      readReqst tableName = readReq tableName
+        selectQuery = readRequestToQuery schema False
-      selectQuery tableName = readRequestToQuery schema False <$> readReqst tableName
+        countQuery = readRequestToCountQuery schema
-      countQuery tableName = readRequestToCountQuery schema <$> readReqst tableName
+        readSqlParts tableName =
-      readSqlParts tableName = (,) <$> selectQuery tableName <*> countQuery tableName
+          let
-      mutationRequest s t = mutateRequest apiRequest t cols (tablePKCols dbStructure s t) =<< fldNames t
+            readReq = readRequest schema tableName maxRows (dbRelations dbStructure) apiRequest
+          in
+          (,,) <$>
+          (selectQuery <$> readReq) <*>
+          (countQuery <$> readReq) <*>
+          (binaryField contentType rawContentTypes returnsScalar =<< readReq)
         mutateSqlParts s t =
-        (,) <$> selectQuery t
+          let
-            <*> (mutateRequestToQuery schema <$> mutationRequest s t)
+            readReq = readRequest s t maxRows (dbRelations dbStructure) apiRequest
-      rawContentTypes =
+            mutReq = mutateRequest apiRequest t cols (tablePKCols dbStructure s t) =<< readReq
-        (decodeContentType <$> configRawMediaTypes conf) `L.union`
+          in
-        [ CTOctetStream, CTTextPlain ]
+          (,) <$>
+          (selectQuery <$> readReq) <*>
+          (mutateRequestToQuery s <$> mutReq)
 
 responseContentTypeOrError :: [ContentType] -> [ContentType] -> Action -> Target -> Either Response ContentType
 responseContentTypeOrError accepts rawContentTypes action target = serves contentTypesForRequest accepts
@@ -375,14 +376,17 @@ responseContentTypeOrError accepts rawContentTypes action target = serves conten
   | If raw(binary) output is requested, check that ContentType is one of the admitted rawContentTypes and that
   | `?select=...` contains only one field other than `*`
 -}
-binaryField :: ContentType -> [ContentType]-> [FieldName] -> Either Response (Maybe FieldName)
+binaryField :: ContentType -> [ContentType] -> Bool -> ReadRequest -> Either Response (Maybe FieldName)
-binaryField ct rawContentTypes fldNames
+binaryField ct rawContentTypes isScalarProc readReq
+  | isScalarProc = Right Nothing
   | ct `elem` rawContentTypes =
       let fieldName = headMay fldNames in
       if length fldNames == 1 && fieldName /= Just "*"
         then Right fieldName
         else Left . errorResponseFor $ BinaryFieldError ct
   | otherwise = Right Nothing
+  where
+    fldNames = fstFieldNames readReq
 
 locationH :: TableName -> [BS.ByteString] -> Header
 locationH tName fields =

src/PostgREST/DbRequestBuilder.hs

@@ -47,8 +47,9 @@ readRequest schema rootTableName maxRows allRels apiRequest  =
   treeRestrictRange maxRows =<<
   augumentRequestWithJoin schema rootRels =<<
   addFiltersOrdersRanges apiRequest <*>
-  (initReadRequest rootName <$> pRequestSelect (iSelect apiRequest))
+  (initReadRequest rootName <$> pRequestSelect sel)
   where
+    sel = fromMaybe "*" $ iSelect apiRequest -- default to all columns requested (SELECT *) for a non existent ?select querystring param
     (rootName, rootRels) = rootWithRelations rootTableName allRels (iAction apiRequest)
 
 -- Get the root table name with its relations according to the Action type.
@@ -296,8 +297,8 @@ addProperty f (targetNodeName:remainingPath, a) (Node rn forest) =
   where
     pathNode = find (\(Node (_,(nodeName,_,alias,_,_)) _) -> nodeName == targetNodeName || alias == Just targetNodeName) forest
 
-mutateRequest :: ApiRequest -> TableName -> S.Set FieldName -> [FieldName] -> [FieldName] -> Either Response MutateRequest
+mutateRequest :: ApiRequest -> TableName -> S.Set FieldName -> [FieldName] -> ReadRequest -> Either Response MutateRequest
-mutateRequest apiRequest tName cols pkCols fldNames = mapLeft errorResponseFor $
+mutateRequest apiRequest tName cols pkCols readReq = mapLeft errorResponseFor $
   case action of
     ActionCreate -> Right $ Insert tName cols ((,) <$> iPreferResolution apiRequest <*> Just pkCols) [] returnings
     ActionUpdate -> Update tName cols <$> combinedLogic <*> pure returnings
@@ -316,7 +317,10 @@ mutateRequest apiRequest tName cols pkCols fldNames = mapLeft errorResponseFor $
     _            -> Left UnsupportedVerb
   where
     action = iAction apiRequest
-    returnings = if iPreferRepresentation apiRequest == None then [] else fldNames
+    returnings =
+      if iPreferRepresentation apiRequest == None
+        then []
+        else returningCols readReq
     filters = map snd <$> mapM pRequestFilter mutateFilters
     logic = map snd <$> mapM pRequestLogicTree logicFilters
     combinedLogic = foldr addFilterToLogicForest <$> logic <*> filters
@@ -324,6 +328,18 @@ mutateRequest apiRequest tName cols pkCols fldNames = mapLeft errorResponseFor $
     (mutateFilters, logicFilters) = join (***) onlyRoot (iFilters apiRequest, iLogic apiRequest)
     onlyRoot = filter (not . ( "." `isInfixOf` ) . fst)
 
+returningCols :: ReadRequest -> [FieldName]
+returningCols rr@(Node _ forest) = fstFieldNames rr ++ (colName <$> fkCols)
+  where
+    -- Without fkCols, when a mutateRequest to /projects?select=name,clients(name) occurs, the RETURNING SQL part would be
+    -- `RETURNING name`(see QueryBuilder).
+    -- This would make the embedding fail because the following JOIN would need the "client_id" column from projects.
+    -- So this adds the foreign key columns to ensure the embedding succeeds, result would be `RETURNING name, client_id`.
+    fkCols = concat $ mapMaybe (\case
+        Node (_, (_, Just Relation{relFColumns=cols, relType=Parent}, _, _, _)) _ -> Just cols
+        _ -> Nothing
+      ) forest
+
 -- Traditional filters(e.g. id=eq.1) are added as root nodes of the LogicTree
 -- they are later concatenated with AND in the QueryBuilder
 addFilterToLogicForest :: Filter -> [LogicTree] -> [LogicTree]

src/PostgREST/Types.hs

@@ -434,13 +434,10 @@ type MutateRequest = MutateQuery
 type ReadNode = (ReadQuery, (NodeName, Maybe Relation, Maybe Alias, Maybe RelationDetail, Depth))
 type Depth = Integer
 
-fieldNames :: ReadRequest -> [FieldName]
+-- First level FieldNames(e.g get a,b from /table?select=a,b,other(c,d))
-fieldNames (Node (sel, _) forest) =
+fstFieldNames :: ReadRequest -> [FieldName]
-  map (fst . view _1) (select sel) ++ map colName fks
+fstFieldNames (Node (sel, _) _) =
-  where
+  fst . view _1 <$> select sel
-    fks = concatMap (fromMaybe [] . f) forest
-    f (Node (_, (_, Just Relation{relFColumns=cols, relType=Parent}, _, _, _)) _) = Just cols
-    f _ = Nothing
 
 data PgVersion = PgVersion {
   pgvNum  :: Int32