refactor: add ResultAggregate type

2023-07-10 15:05:51 -05:00
parent 774d015eb5
commit 905fcb05cc
6 changed files with 90 additions and 64 deletions
@@ -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
@@ -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
@@ -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)
@@ -7,12 +7,7 @@ Description : Helper functions for PostgREST.QueryBuilder.
 -}
 module PostgREST.Query.SqlFragment
  ( noLocationF
-  , asBinaryF
-  , asCsvF
-  , asGeoJsonF
-  , asJsonF
-  , asJsonSingleF
-  , asXmlF
+  , aggF
  , 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 fieldName = "coalesce(xmlagg(_postgrest_t." <> pgFmtIdent fieldName <> "), '')"
+asXmlF :: Maybe FieldName -> SQL.Snippet
+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 fieldName = "coalesce(string_agg(_postgrest_t." <> pgFmtIdent fieldName <> ", ''), '')"
+asBinaryF :: Maybe FieldName -> SQL.Snippet
+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"
@@ -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 (..),
-                                          MediaType (..),
-                                          getMediaType)
+import PostgREST.MediaType              (MTPlanFormat (..),
+                                         MediaType (..))
 import PostgREST.Query.SqlFragment
-import PostgREST.SchemaCache.Identifiers (FieldName)
-import PostgREST.SchemaCache.Routine     (Routine)
+import PostgREST.SchemaCache.Routine    (ResultAggregate (..),
+                                         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
-    | otherwise   = selectQuery
+    | rAgg == NoAgg = "SELECT * FROM " <> sourceCTE
+    | 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 selectQuery countQuery countTotal mt binaryField =
+prepareRead :: SQL.Snippet -> SQL.Snippet -> Bool -> MediaType -> ResultAggregate -> Bool -> SQL.Statement () ResultSet
+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
@@ -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