refactor: add ResultAggregate type

src/PostgREST/MediaType.hs

@@ -30,6 +30,7 @@ data MediaType
   | MTOctetStream
   | MTAny
   | MTOther ByteString
+  -- TODO MTPlan should only have its options as [Text]. Its ResultAggregate should have the typed attributes.
   | MTPlan MediaType MTPlanFormat [MTPlanOption]
   deriving Show
 instance Eq MediaType where

src/PostgREST/Plan.hs

@@ -59,7 +59,8 @@ import PostgREST.SchemaCache.Relationship    (Cardinality (..),
                                               relIsToOne)
 import PostgREST.SchemaCache.Representations (DataRepresentation (..),
                                               RepresentationsMap)
-import PostgREST.SchemaCache.Routine         (Routine (..),
+import PostgREST.SchemaCache.Routine         (ResultAggregate (..),
+                                              Routine (..),
                                               RoutineMap,
                                               RoutineParam (..),
                                               funcReturnsCompositeAlias,
@@ -89,13 +90,14 @@ import Protolude hiding (from)
 data WrappedReadPlan = WrappedReadPlan {
   wrReadPlan :: ReadPlanTree
 , wrTxMode   :: SQL.Mode
-, wrBinField :: Maybe FieldName
+, wrResAgg   :: ResultAggregate
 }
 
 data MutateReadPlan = MutateReadPlan {
   mrReadPlan   :: ReadPlanTree
 , mrMutatePlan :: MutatePlan
 , mrTxMode     :: SQL.Mode
+, mrResAgg     :: ResultAggregate
 }
 
 data CallReadPlan = CallReadPlan {
@@ -103,20 +105,21 @@ data CallReadPlan = CallReadPlan {
 , crCallPlan :: CallPlan
 , crTxMode   :: SQL.Mode
 , crProc     :: Routine
-, crBinField :: Maybe FieldName
+, crResAgg   :: ResultAggregate
 }
 
 wrappedReadPlan :: QualifiedIdentifier -> AppConfig -> SchemaCache -> ApiRequest -> Either Error WrappedReadPlan
 wrappedReadPlan  identifier conf sCache apiRequest = do
   rPlan <- readPlan identifier conf sCache apiRequest
   binField <- mapLeft ApiRequestError $ binaryField conf (iAcceptMediaType apiRequest) Nothing rPlan
-  return $ WrappedReadPlan rPlan SQL.Read binField
+  return $ WrappedReadPlan rPlan SQL.Read $ mediaToAggregate (iAcceptMediaType apiRequest) binField Nothing
 
 mutateReadPlan :: Mutation -> ApiRequest -> QualifiedIdentifier -> AppConfig -> SchemaCache -> Either Error MutateReadPlan
-mutateReadPlan  mutation apiRequest identifier conf sCache = do
+mutateReadPlan  mutation apiRequest@ApiRequest{iPreferences=Preferences{preferRepresentation}} identifier conf sCache = do
   rPlan <- readPlan identifier conf sCache apiRequest
+  binField <- mapLeft ApiRequestError $ binaryField conf (iAcceptMediaType apiRequest) Nothing rPlan
   mPlan <- mutatePlan mutation identifier apiRequest sCache rPlan
-  return $ MutateReadPlan rPlan mPlan SQL.Write
+  return $ MutateReadPlan rPlan mPlan SQL.Write $ mediaToAggregate (iAcceptMediaType apiRequest) binField (Just preferRepresentation)
 
 callReadPlan :: QualifiedIdentifier -> AppConfig -> SchemaCache -> ApiRequest -> InvokeMethod -> Either Error CallReadPlan
 callReadPlan identifier conf sCache apiRequest invMethod = do
@@ -141,7 +144,7 @@ callReadPlan identifier conf sCache apiRequest invMethod = do
           (InvPost, Routine.Volatile)  -> SQL.Write
       cPlan = callPlan proc apiRequest paramKeys args rPlan
   binField <- mapLeft ApiRequestError $ binaryField conf (iAcceptMediaType apiRequest) (Just proc) rPlan
-  return $ CallReadPlan rPlan cPlan txMode proc binField
+  return $ CallReadPlan rPlan cPlan txMode proc $ mediaToAggregate (iAcceptMediaType apiRequest) binField Nothing
   where
     Preferences{..} = iPreferences apiRequest
     qsParams' = QueryParams.qsParams (iQueryParams apiRequest)
@@ -835,3 +838,27 @@ binaryField AppConfig{configRawMediaTypes} acceptMediaType proc rpTree
     fstFieldName (Node ReadPlan{select=(CoercibleField{cfName="*", cfJsonPath=[]}, _, _):_} [])  = Nothing
     fstFieldName (Node ReadPlan{select=[(CoercibleField{cfName=fld, cfJsonPath=[]}, _, _)]} [])  = Just fld
     fstFieldName _                                               = Nothing
+
+mediaToAggregate :: MediaType -> Maybe FieldName -> Maybe PreferRepresentation ->  ResultAggregate
+mediaToAggregate mt binField rep =
+  if rep == Just HeadersOnly || rep == Just None
+    then NoAgg
+  else case mt of
+    MTApplicationJSON     -> BuiltinAggJson
+    MTSingularJSON        -> BuiltinAggSingleJson
+    MTGeoJSON             -> BuiltinAggGeoJson
+    MTTextCSV             -> BuiltinAggCsv
+    MTAny                 -> BuiltinAggJson
+    MTOpenAPI             -> BuiltinAggJson
+    MTUrlEncoded          -> NoAgg -- TODO: unreachable since a previous step (producedMediaTypes) whitelists the media types that can become aggregates.
+
+    -- binary types
+    MTTextPlain           -> BuiltinAggBinary binField
+    MTTextXML             -> BuiltinAggXml binField
+    MTOctetStream         -> BuiltinAggBinary binField
+    MTOther _             -> BuiltinAggBinary binField
+
+    -- Doing `Accept: application/vnd.pgrst.plan; for="application/vnd.pgrst.plan"` doesn't make sense, so we just empty the body.
+    -- TODO: fail instead to be more strict
+    MTPlan (MTPlan{}) _ _ -> NoAgg
+    MTPlan media      _ _ -> mediaToAggregate media binField rep

src/PostgREST/Query.hs

@@ -65,7 +65,7 @@ import Protolude hiding (Handler)
 type DbHandler = ExceptT Error SQL.Transaction
 
 readQuery :: WrappedReadPlan -> AppConfig -> ApiRequest -> DbHandler ResultSet
-readQuery WrappedReadPlan{wrReadPlan, wrBinField} conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} = do
+readQuery WrappedReadPlan{wrReadPlan, wrResAgg} conf@AppConfig{..} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} = do
   let countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan
   resultSet <-
      lift . SQL.statement mempty $
@@ -79,7 +79,7 @@ readQuery WrappedReadPlan{wrReadPlan, wrBinField} conf@AppConfig{..} apiReq@ApiR
         )
         (shouldCount preferCount)
         iAcceptMediaType
-        wrBinField
+        wrResAgg
         configDbPreparedStatements
   failNotSingular iAcceptMediaType resultSet
   optionalRollback conf apiReq
@@ -150,7 +150,7 @@ deleteQuery mrPlan apiReq@ApiRequest{..} conf = do
   pure resultSet
 
 invokeQuery :: Routine -> CallReadPlan -> ApiRequest -> AppConfig -> PgVersion -> DbHandler ResultSet
-invokeQuery rout CallReadPlan{crReadPlan, crCallPlan, crBinField} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} conf@AppConfig{..} pgVer = do
+invokeQuery rout CallReadPlan{crReadPlan, crCallPlan, crResAgg} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} conf@AppConfig{..} pgVer = do
   resultSet <-
     lift . SQL.statement mempty $
       Statements.prepareCall
@@ -160,7 +160,7 @@ invokeQuery rout CallReadPlan{crReadPlan, crCallPlan, crBinField} apiReq@ApiRequ
         (QueryBuilder.readPlanToCountQuery crReadPlan)
         (shouldCount preferCount)
         iAcceptMediaType
-        crBinField
+        crResAgg
         configDbPreparedStatements
 
   optionalRollback conf apiReq
@@ -185,7 +185,7 @@ openApiQuery sCache pgVer AppConfig{..} tSchema =
       pure Nothing
 
 writeQuery :: MutateReadPlan -> ApiRequest -> AppConfig  -> DbHandler ResultSet
-writeQuery MutateReadPlan{mrReadPlan, mrMutatePlan} apiReq@ApiRequest{iPreferences=Preferences{..}} conf =
+writeQuery MutateReadPlan{mrReadPlan, mrMutatePlan, mrResAgg} apiReq@ApiRequest{iPreferences=Preferences{..}} conf =
   let
     (isInsert, pkCols) = case mrMutatePlan of {Insert{insPkCols} -> (True, insPkCols); _ -> (False, mempty);}
   in
@@ -195,6 +195,7 @@ writeQuery MutateReadPlan{mrReadPlan, mrMutatePlan} apiReq@ApiRequest{iPreferenc
       (QueryBuilder.mutatePlanToQuery mrMutatePlan)
       isInsert
       (iAcceptMediaType apiReq)
+      mrResAgg
       preferRepresentation
       pkCols
       (configDbPreparedStatements conf)

src/PostgREST/Query/SqlFragment.hs

@@ -7,12 +7,7 @@ Description : Helper functions for PostgREST.QueryBuilder.
 -}
 module PostgREST.Query.SqlFragment
   ( noLocationF
-  , asBinaryF
+  , aggF
-  , asCsvF
-  , asGeoJsonF
-  , asJsonF
-  , asJsonSingleF
-  , asXmlF
   , countF
   , fromQi
   , limitOffsetF
@@ -81,7 +76,8 @@ import PostgREST.RangeQuery              (NonnegRange, allRange,
                                           rangeLimit, rangeOffset)
 import PostgREST.SchemaCache.Identifiers (FieldName,
                                           QualifiedIdentifier (..))
-import PostgREST.SchemaCache.Routine     (Routine (..),
+import PostgREST.SchemaCache.Routine     (ResultAggregate (..),
+                                          Routine (..),
                                           funcReturnsScalar,
                                           funcReturnsSetOfScalar,
                                           funcReturnsSingleComposite)
@@ -208,14 +204,18 @@ asJsonF rout
       Just r  -> (funcReturnsSingleComposite r, funcReturnsScalar r, funcReturnsSetOfScalar r)
       Nothing -> (False, False, False)
 
-asXmlF :: FieldName -> SQL.Snippet
+asXmlF :: Maybe FieldName -> SQL.Snippet
-asXmlF fieldName = "coalesce(xmlagg(_postgrest_t." <> pgFmtIdent fieldName <> "), '')"
+asXmlF (Just fieldName) = "coalesce(xmlagg(_postgrest_t." <> pgFmtIdent fieldName <> "), '')"
+-- TODO unreachable because a previous step(binaryField) will validate that there's a field. This will be cleared once custom media types are implemented.
+asXmlF Nothing          = "coalesce(xmlagg(_postgrest_t), '')"
 
 asGeoJsonF ::  SQL.Snippet
 asGeoJsonF = "json_build_object('type', 'FeatureCollection', 'features', coalesce(json_agg(ST_AsGeoJSON(_postgrest_t)::json), '[]'))"
 
-asBinaryF :: FieldName -> SQL.Snippet
+asBinaryF :: Maybe FieldName -> SQL.Snippet
-asBinaryF fieldName = "coalesce(string_agg(_postgrest_t." <> pgFmtIdent fieldName <> ", ''), '')"
+asBinaryF (Just fieldName) = "coalesce(string_agg(_postgrest_t." <> pgFmtIdent fieldName <> ", ''), '')"
+-- TODO unreachable because a previous step(binaryField) will validate that there's a field. This will be cleared once custom media types are implemented.
+asBinaryF Nothing          = "coalesce(string_agg(_postgrest_t, ''), '')"
 
 locationF :: [Text] -> SQL.Snippet
 locationF pKeys = [qc|(
@@ -491,3 +491,13 @@ setConfigLocalJson prefix keyVals = [setConfigLocal mempty (prefix, gucJsonVal k
     gucJsonVal = LBS.toStrict . JSON.encode . HM.fromList . arrayByteStringToText
     arrayByteStringToText :: [(ByteString, ByteString)] -> [(Text,Text)]
     arrayByteStringToText keyVal = (T.decodeUtf8 *** T.decodeUtf8) <$> keyVal
+
+aggF :: Maybe Routine -> ResultAggregate -> SQL.Snippet
+aggF rout = \case
+  BuiltinAggJson          -> asJsonF rout
+  BuiltinAggSingleJson    -> asJsonSingleF rout
+  BuiltinAggGeoJson       -> asGeoJsonF
+  BuiltinAggCsv           -> asCsvF
+  BuiltinAggXml    bField -> asXmlF    bField
+  BuiltinAggBinary bField -> asBinaryF bField
+  NoAgg                   -> "''::text"

src/PostgREST/Query/Statements.hs

@@ -23,15 +23,13 @@ import qualified Hasql.DynamicStatements.Statement as SQL
 import qualified Hasql.Statement                   as SQL
 
 import Control.Lens ((^?))
-import Data.Maybe   (fromJust)
 
 import PostgREST.ApiRequest.Preferences
-import PostgREST.MediaType               (MTPlanFormat (..),
+import PostgREST.MediaType              (MTPlanFormat (..),
-                                          MediaType (..),
+                                         MediaType (..))
-                                          getMediaType)
 import PostgREST.Query.SqlFragment
-import PostgREST.SchemaCache.Identifiers (FieldName)
+import PostgREST.SchemaCache.Routine    (ResultAggregate (..),
-import PostgREST.SchemaCache.Routine     (Routine)
+                                         Routine)
 
 import Protolude
 
@@ -55,9 +53,9 @@ data ResultSet
   | RSPlan BS.ByteString -- ^ the plan of the query
 
 
-prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType ->
+prepareWrite :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> ResultAggregate ->
                 PreferRepresentation -> [Text] -> Bool -> SQL.Statement () ResultSet
-prepareWrite selectQuery mutateQuery isInsert mt rep pKeys =
+prepareWrite selectQuery mutateQuery isInsert mt rAgg rep pKeys =
   SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
  where
   snippet =
@@ -66,7 +64,7 @@ prepareWrite selectQuery mutateQuery isInsert mt rep pKeys =
       "'' AS total_result_set, " <>
       "pg_catalog.count(_postgrest_t) AS page_total, " <>
       locF <> " AS header, " <>
-      bodyF <> " AS body, " <>
+      aggF Nothing rAgg <> " AS body, " <>
       responseHeadersF <> " AS response_headers, " <>
       responseStatusF  <> " AS response_status " <>
     "FROM (" <> selectF <> ") _postgrest_t"
@@ -80,25 +78,18 @@ prepareWrite selectQuery mutateQuery isInsert mt rep pKeys =
         "END"
       else noLocationF
 
-  bodyF
-    | rep /= Full                       = "''"
-    | getMediaType mt == MTTextCSV      = asCsvF
-    | getMediaType mt == MTGeoJSON      = asGeoJsonF
-    | getMediaType mt == MTSingularJSON = asJsonSingleF Nothing
-    | otherwise                         = asJsonF Nothing
-
   selectF
     -- prevent using any of the column names in ?select= when no response is returned from the CTE
-    | rep /= Full = "SELECT * FROM " <> sourceCTE
+    | rAgg == NoAgg = "SELECT * FROM " <> sourceCTE
-    | otherwise   = selectQuery
+    | otherwise     = selectQuery
 
   decodeIt :: HD.Result ResultSet
   decodeIt = case mt of
     MTPlan{} -> planRow
     _        -> fromMaybe (RSStandard Nothing 0 mempty mempty Nothing Nothing) <$> HD.rowMaybe (standardRow False)
 
-prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> Maybe FieldName -> Bool -> SQL.Statement () ResultSet
+prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> ResultAggregate -> Bool -> SQL.Statement () ResultSet
-prepareRead selectQuery countQuery countTotal mt binaryField =
+prepareRead selectQuery countQuery countTotal mt rAgg =
   SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
  where
   snippet =
@@ -107,30 +98,22 @@ prepareRead selectQuery countQuery countTotal mt binaryField =
     "SELECT " <>
       countResultF <> " AS total_result_set, " <>
       "pg_catalog.count(_postgrest_t) AS page_total, " <>
-      bodyF <> " AS body, " <>
+      aggF Nothing rAgg <> " AS body, " <>
       responseHeadersF <> " AS response_headers, " <>
       responseStatusF <> " AS response_status " <>
     "FROM ( SELECT * FROM " <> sourceCTE <> " ) _postgrest_t"
 
   (countCTEF, countResultF) = countF countQuery countTotal
 
-  bodyF
-    | getMediaType mt == MTTextCSV                       = asCsvF
-    | getMediaType mt == MTSingularJSON                  = asJsonSingleF Nothing
-    | getMediaType mt == MTGeoJSON                       = asGeoJsonF
-    | isJust binaryField && getMediaType mt == MTTextXML = asXmlF $ fromJust binaryField
-    | isJust binaryField                                 = asBinaryF $ fromJust binaryField
-    | otherwise                                          = asJsonF Nothing
-
   decodeIt :: HD.Result ResultSet
   decodeIt = case mt of
     MTPlan{} -> planRow
     _        -> HD.singleRow $ standardRow True
 
 prepareCall :: Routine -> SQL.Snippet -> SQL.Snippet -> SQL.Snippet -> Bool ->
-               MediaType -> Maybe FieldName -> Bool ->
+               MediaType -> ResultAggregate -> Bool ->
                SQL.Statement () ResultSet
-prepareCall rout callProcQuery selectQuery countQuery countTotal mt binaryField =
+prepareCall rout callProcQuery selectQuery countQuery countTotal mt rAgg =
   SQL.dynamicallyParameterized (mtSnippet mt snippet) decodeIt
   where
     snippet =
@@ -139,21 +122,13 @@ prepareCall rout callProcQuery selectQuery countQuery countTotal mt binaryField
       "SELECT " <>
         countResultF <> " AS total_result_set, " <>
         "pg_catalog.count(_postgrest_t) AS page_total, " <>
-        bodyF <> " AS body, " <>
+        aggF (Just rout) rAgg <> " AS body, " <>
         responseHeadersF <> " AS response_headers, " <>
         responseStatusF <> " AS response_status " <>
       "FROM (" <> selectQuery <> ") _postgrest_t"
 
     (countCTEF, countResultF) = countF countQuery countTotal
 
-    bodyF
-     | getMediaType mt == MTSingularJSON                  = asJsonSingleF $ Just rout
-     | getMediaType mt == MTTextCSV                       = asCsvF
-     | getMediaType mt == MTGeoJSON                       = asGeoJsonF
-     | isJust binaryField && getMediaType mt == MTTextXML = asXmlF $ fromJust binaryField
-     | isJust binaryField                                 = asBinaryF $ fromJust binaryField
-     | otherwise                                          = asJsonF $ Just rout
-
     decodeIt :: HD.Result ResultSet
     decodeIt = case mt of
       MTPlan{} -> planRow

src/PostgREST/SchemaCache/Routine.hs

@@ -14,6 +14,7 @@ module PostgREST.SchemaCache.Routine
   , funcReturnsVoid
   , funcTableName
   , funcReturnsCompositeAlias
+  , ResultAggregate(..)
   ) where
 
 import           Data.Aeson                 ((.=))
@@ -21,7 +22,8 @@ import qualified Data.Aeson                 as JSON
 import qualified Data.HashMap.Strict        as HM
 import qualified Hasql.Transaction.Sessions as SQL
 
-import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..),
+import PostgREST.SchemaCache.Identifiers (FieldName,
+                                          QualifiedIdentifier (..),
                                           Schema, TableName)
 
 import Protolude
@@ -85,6 +87,16 @@ instance Ord Routine where
 -- | It uses a HashMap for a faster lookup.
 type RoutineMap = HM.HashMap QualifiedIdentifier [Routine]
 
+data ResultAggregate
+   = BuiltinAggJson
+   | BuiltinAggSingleJson
+   | BuiltinAggGeoJson
+   | BuiltinAggCsv
+   | BuiltinAggXml    (Maybe FieldName)
+   | BuiltinAggBinary (Maybe FieldName)
+   | NoAgg
+   deriving (Eq, Show)
+
 funcReturnsScalar :: Routine -> Bool
 funcReturnsScalar proc = case proc of
   Function{pdReturnType = Single (Scalar{})} -> True