refactor: remove fromJust in DbRequestBuilder

* Also comment and reorganize DbRequestBuilder

src/PostgREST/App.hs

@@ -48,8 +48,7 @@ import PostgREST.ApiRequest       (Action (..), ApiRequest (..),
 import PostgREST.Auth             (containsRole, jwtClaims,
                                    parseSecret)
 import PostgREST.Config           (AppConfig (..))
-import PostgREST.DbRequestBuilder (fieldNames, mutateRequest,
+import PostgREST.DbRequestBuilder (mutateRequest, readRequest)
-                                   readRequest)
 import PostgREST.DbStructure
 import PostgREST.Error            (PgError (..), SimpleError (..),
                                    errorResponseFor, singularityError)
@@ -125,9 +124,9 @@ app dbStructure proc cols conf apiRequest =
     Right contentType ->
       case (iAction apiRequest, iTarget apiRequest, iPayload apiRequest) of
 
-        (ActionRead headersOnly, TargetIdent qi, Nothing) ->
+        (ActionRead headersOnly, TargetIdent (QualifiedIdentifier _ tName), Nothing) ->
-          let partsField = (,) <$> readSqlParts
+          let partsField = (,) <$> readSqlParts tName
-                <*> (binaryField contentType rawContentTypes =<< fldNames) in
+                <*> (binaryField contentType rawContentTypes =<< fldNames tName) in
           case partsField of
             Left errorResponse -> return errorResponse
             Right ((q, cq), bField) -> do
@@ -151,7 +150,7 @@ app dbStructure proc cols conf apiRequest =
                   then errorResponseFor . singularityError $ queryTotal
                   else responseLBS status
                     [toHeader contentType, contentRange,
-                     contentLocationH (qiName qi) (iCanonicalQS apiRequest)]
+                     contentLocationH tName (iCanonicalQS apiRequest)]
                     (if headersOnly then mempty else toS body)
 
         (ActionCreate, TargetIdent (QualifiedIdentifier tSchema tName), Just pJson) ->
@@ -283,12 +282,13 @@ app dbStructure proc cols conf apiRequest =
                   allOrigins = ("Access-Control-Allow-Origin", "*") :: Header in
               return $ responseLBS status200 [allOrigins, allowH] mempty
 
-        (ActionInvoke invMethod, TargetProc qi _, Just pJson) ->
+        (ActionInvoke invMethod, TargetProc qi@(QualifiedIdentifier _ pName) _, Just pJson) ->
           let returnsScalar = maybe False procReturnsScalar proc
+              tName = fromMaybe pName $ procTableName =<< proc
               rpcBinaryField = if returnsScalar
                                  then Right Nothing
-                                 else binaryField contentType rawContentTypes =<< fldNames
+                                 else binaryField contentType rawContentTypes =<< fldNames tName
-              parts = (,) <$> readSqlParts <*> rpcBinaryField in
+              parts = (,) <$> readSqlParts tName <*> rpcBinaryField in
           case parts of
             Left errorResponse -> return errorResponse
             Right ((q, cq), bField) -> do
@@ -337,15 +337,15 @@ app dbStructure proc cols conf apiRequest =
       plannedCount = iPreferCount apiRequest == Just PlannedCount
       shouldCount = exactCount || estimatedCount
       topLevelRange = iTopLevelRange apiRequest
-      readReq = readRequest maxRows (dbRelations dbStructure) proc apiRequest
+      readReq tableName = readRequest schema tableName maxRows (dbRelations dbStructure) apiRequest
-      fldNames = fieldNames <$> readReq
+      fldNames tableName = fieldNames <$> readReq tableName
-      readDbRequest = DbRead <$> readReq
+      readDbRequest tableName = DbRead <$> readReq tableName
-      selectQuery = requestToQuery schema False <$> readDbRequest
+      selectQuery tableName = requestToQuery schema False <$> readDbRequest tableName
-      countQuery = requestToCountQuery schema <$> readDbRequest
+      countQuery tableName = requestToCountQuery schema <$> readDbRequest tableName
-      readSqlParts = (,) <$> selectQuery <*> countQuery
+      readSqlParts tableName = (,) <$> selectQuery tableName <*> countQuery tableName
-      mutationDbRequest s t = mutateRequest apiRequest t cols (tablePKCols dbStructure s t) =<< fldNames
+      mutationDbRequest s t = mutateRequest apiRequest t cols (tablePKCols dbStructure s t) =<< fldNames t
       mutateSqlParts s t =
-        (,) <$> selectQuery
+        (,) <$> selectQuery t
             <*> (requestToQuery schema False . DbMutate <$> mutationDbRequest s t)
       rawContentTypes =
         (decodeContentType <$> configRawMediaTypes conf) `L.union`

src/PostgREST/DbRequestBuilder.hs

@@ -14,7 +14,6 @@ A query tree is built in case of resource embedding. By inferring the relationsh
 module PostgREST.DbRequestBuilder (
   readRequest
 , mutateRequest
-, fieldNames
 ) where
 
 import qualified Data.ByteString.Char8 as BS
@@ -22,8 +21,6 @@ import qualified Data.HashMap.Strict   as M
 import qualified Data.Set              as S
 
 import Control.Arrow           ((***))
-import Control.Lens.Getter     (view)
-import Control.Lens.Tuple      (_1)
 import Data.Either.Combinators (mapLeft)
 import Data.Foldable           (foldr1)
 import Data.List               (delete)
@@ -37,77 +34,61 @@ import Data.Tree
 import Network.Wai
 
 import PostgREST.ApiRequest (Action (..), ApiRequest (..),
-                             PreferRepresentation (..),
+                             PreferRepresentation (..))
-                             PreferRepresentation (..), Target (..))
 import PostgREST.Error      (ApiRequestError (..), errorResponseFor)
 import PostgREST.Parsers
 import PostgREST.RangeQuery (NonnegRange, allRange, restrictRange)
 import PostgREST.Types
 import Protolude            hiding (from)
 
-readRequest :: Maybe Integer -> [Relation] -> Maybe ProcDescription -> ApiRequest -> Either Response ReadRequest
+readRequest :: Schema -> TableName -> Maybe Integer -> [Relation] -> ApiRequest -> Either Response ReadRequest
-readRequest maxRows allRels proc apiRequest  =
+readRequest schema rootTableName maxRows allRels apiRequest  =
   mapLeft errorResponseFor $
   treeRestrictRange maxRows =<<
-  augumentRequestWithJoin schema relations =<<
+  augumentRequestWithJoin schema rootRels =<<
   addFiltersOrdersRanges apiRequest <*>
-  (buildReadRequest <$> pRequestSelect (iSelect apiRequest))
+  (initReadRequest rootName <$> pRequestSelect (iSelect apiRequest))
   where
-    action = iAction apiRequest
+    (rootName, rootRels) = rootWithRelations rootTableName allRels (iAction apiRequest)
-    (schema, rootTableName) = fromJust $ -- Make it safe
-      let target = iTarget apiRequest in
-      case target of
-        (TargetIdent (QualifiedIdentifier s t) ) -> Just (s, t)
-        (TargetProc  (QualifiedIdentifier s pName) _ ) -> Just (s, tName)
-          where
-            tName = case pdReturnType <$> proc of
-              Just (SetOf (Composite qi))  -> qiName qi
-              Just (Single (Composite qi)) -> qiName qi
-              _                            -> pName
 
-        _ -> Nothing
+-- Get the root table name with its relations according to the Action type.
+-- This is done because of the shape of the final SQL Query. The mutation cases are wrapped in a WITH {sourceCTEName}(see Statements.hs).
+-- So we need a FROM {sourceCTEName} instead of FROM {tableName}.
+rootWithRelations :: TableName -> [Relation] -> Action -> (TableName, [Relation])
+rootWithRelations rootTableName allRels action = case action of
+  ActionRead _ -> (rootTableName, allRels) -- normal read case
+  _            -> (sourceCTEName, mapMaybe toSourceRelation allRels ++ allRels) -- mutation cases and calling proc
+  where
+    -- in a relation where one of the tables matches "TableName"
+    -- replace the name to that table with pg_source
+    -- this "fake" relations is needed so that in a mutate query or proc call
+    -- we can look at the "returning *" part which is wrapped with a "with pg_source"
+    -- as just another table that has relations with other tables
+    toSourceRelation :: Relation -> Maybe Relation
+    toSourceRelation r@(Relation t _ ft _ _ rt _ _)
+      | rootTableName == tableName t = Just $ r {relTable=t {tableName=sourceCTEName}}
+      | rootTableName == tableName ft = Just $ r {relFTable=t {tableName=sourceCTEName}}
+      | Just rootTableName == (tableName <$> rt) = Just $ r {relLinkTable=(\tbl -> tbl {tableName=sourceCTEName}) <$> rt}
+      | otherwise = Nothing
 
-    -- Build tree with a Depth attribute so when a self join occurs we can differentiate the parent and child tables by having
+-- Build the initial tree with a Depth attribute so when a self join occurs we can differentiate the parent and child tables by having
-    -- an alias like "table_depth", this is related to issue #987.
+-- an alias like "table_depth", this is related to http://github.com/PostgREST/postgrest/issues/987.
-    buildReadRequest :: [Tree SelectItem] -> ReadRequest
+initReadRequest :: TableName -> [Tree SelectItem] -> ReadRequest
-    buildReadRequest fieldTree =
+initReadRequest rootTableName =
-      let rootDepth = 0
+  foldr (treeEntry rootDepth) initial
-          rootNodeName = case action of
+  where
-            ActionRead _ -> rootTableName
+    rootDepth = 0
-            _            -> sourceCTEName in
+    initial = Node (Select [] rootTableName Nothing [] [] [] [] allRange, (rootTableName, Nothing, Nothing, Nothing, rootDepth)) []
-      foldr (treeEntry rootDepth) (Node (Select [] rootNodeName Nothing [] [] [] [] allRange, (rootNodeName, Nothing, Nothing, Nothing, rootDepth)) []) fieldTree
+    treeEntry :: Depth -> Tree SelectItem -> ReadRequest -> ReadRequest
-      where
+    treeEntry depth (Node fld@((fn, _),_,alias,relationDetail) fldForest) (Node (q, i) rForest) =
-        treeEntry :: Depth -> Tree SelectItem -> ReadRequest -> ReadRequest
+      let nxtDepth = succ depth in
-        treeEntry depth (Node fld@((fn, _),_,alias,relationDetail) fldForest) (Node (q, i) rForest) =
+      case fldForest of
-          let nxtDepth = succ depth in
+        [] -> Node (q {select=fld:select q}, i) rForest
-          case fldForest of
+        _  -> Node (q, i) $
-            [] -> Node (q {select=fld:select q}, i) rForest
+              foldr (treeEntry nxtDepth) (Node (Select [] fn Nothing [] [] [] [] allRange, (fn, Nothing, alias, relationDetail, nxtDepth)) []) fldForest:rForest
-            _  -> Node (q, i) $
-                  foldr (treeEntry nxtDepth) (Node (Select [] fn Nothing [] [] [] [] allRange, (fn, Nothing, alias, relationDetail, nxtDepth)) []) fldForest:rForest
-
-    relations :: [Relation]
-    relations = case action of
-      ActionCreate   -> fakeSourceRelations ++ allRels
-      ActionUpdate   -> fakeSourceRelations ++ allRels
-      ActionDelete   -> fakeSourceRelations ++ allRels
-      ActionInvoke _ -> fakeSourceRelations ++ allRels
-      _              -> allRels
-      where fakeSourceRelations = mapMaybe (toSourceRelation rootTableName) allRels
-
--- in a relation where one of the tables matches "TableName"
--- replace the name to that table with pg_source
--- this "fake" relations is needed so that in a mutate query
--- we can look at the "returning *" part which is wrapped with a "with"
--- as just another table that has relations with other tables
-toSourceRelation :: TableName -> Relation -> Maybe Relation
-toSourceRelation mt r@(Relation t _ ft _ _ rt _ _)
-  | mt == tableName t = Just $ r {relTable=t {tableName=sourceCTEName}}
-  | mt == tableName ft = Just $ r {relFTable=t {tableName=sourceCTEName}}
-  | Just mt == (tableName <$> rt) = Just $ r {relLinkTable=(\tbl -> tbl {tableName=sourceCTEName}) <$> rt}
-  | otherwise = Nothing
 
 treeRestrictRange :: Maybe Integer -> ReadRequest -> Either ApiRequestError ReadRequest
-treeRestrictRange maxRows_ request = pure $ nodeRestrictRange maxRows_ `fmap` request
+treeRestrictRange maxRows request = pure $ nodeRestrictRange maxRows <$> request
   where
     nodeRestrictRange :: Maybe Integer -> ReadNode -> ReadNode
     nodeRestrictRange m (q@Select {range_=r}, i) = (q{range_=restrictRange m r }, i)
@@ -268,7 +249,7 @@ addFiltersOrdersRanges apiRequest = foldr1 (liftA2 (.)) [
   ]
   {-
   The esence of what is going on above is that we are composing tree functions
-  of type (ReadRequest->ReadRequest) that are in (Either ParseError a) context
+  of type (ReadRequest->ReadRequest) that are in (Either ApiRequestError a) context
   -}
   where
     filters :: Either ApiRequestError [(EmbedPath, Filter)]
@@ -348,14 +329,6 @@ mutateRequest apiRequest tName cols pkCols fldNames = mapLeft errorResponseFor $
     (mutateFilters, logicFilters) = join (***) onlyRoot (iFilters apiRequest, iLogic apiRequest)
     onlyRoot = filter (not . ( "." `isInfixOf` ) . fst)
 
-fieldNames :: ReadRequest -> [FieldName]
-fieldNames (Node (sel, _) forest) =
-  map (fst . view _1) (select sel) ++ map colName fks
-  where
-    fks = concatMap (fromMaybe [] . f) forest
-    f (Node (_, (_, Just Relation{relFColumns=cols, relType=Parent}, _, _, _)) _) = Just cols
-    f _ = Nothing
-
 -- 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

@@ -6,6 +6,9 @@ Description : PostgREST common types and functions used by the rest of the modul
 
 module PostgREST.Types where
 
+import Control.Lens.Getter (view)
+import Control.Lens.Tuple  (_1)
+
 import qualified Data.Aeson               as JSON
 import qualified Data.ByteString          as BS
 import qualified Data.ByteString.Internal as BS (c2w)
@@ -166,6 +169,12 @@ procReturnsScalar proc = case proc of
   ProcDescription{pdReturnType = (Single (Scalar _))} -> True
   _                                                   -> False
 
+procTableName :: ProcDescription -> Maybe TableName
+procTableName proc = case pdReturnType proc of
+  SetOf  (Composite qi) -> Just $ qiName qi
+  Single (Composite qi) -> Just $ qiName qi
+  _                     -> Nothing
+
 type Schema = Text
 type TableName = Text
 
@@ -426,6 +435,14 @@ type ReadNode = (ReadQuery, (NodeName, Maybe Relation, Maybe Alias, Maybe Relati
 type Depth = Integer
 type MutateRequest = MutateQuery
 
+fieldNames :: ReadRequest -> [FieldName]
+fieldNames (Node (sel, _) forest) =
+  map (fst . view _1) (select sel) ++ map colName fks
+  where
+    fks = concatMap (fromMaybe [] . f) forest
+    f (Node (_, (_, Just Relation{relFColumns=cols, relType=Parent}, _, _, _)) _) = Just cols
+    f _ = Nothing
+
 data PgVersion = PgVersion {
   pgvNum  :: Int32
 , pgvName :: Text