feat: add max-affected preference to prefer header

CHANGELOG.md

@@ -5,6 +5,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+ - #2887, Add Preference `max-affected` to limit affected resources - @taimoorzaeem
+
 ## [12.0.1] - 2023-12-12
 
 ### Fixed

src/PostgREST/ApiRequest/Preferences.hs

@@ -17,6 +17,7 @@ module PostgREST.ApiRequest.Preferences
   , PreferResolution(..)
   , PreferTransaction(..)
   , PreferTimezone(..)
+  , PreferMaxAffected(..)
   , fromHeaders
   , shouldCount
   , prefAppliedHeader
@@ -42,6 +43,7 @@ import Protolude
 -- >>> deriving instance Show PreferMissing
 -- >>> deriving instance Show PreferHandling
 -- >>> deriving instance Show PreferTimezone
+-- >>> deriving instance Show PreferMaxAffected
 -- >>> deriving instance Show Preferences
 
 -- | Preferences recognized by the application.
@@ -55,6 +57,7 @@ data Preferences
     , preferMissing        :: Maybe PreferMissing
     , preferHandling       :: Maybe PreferHandling
     , preferTimezone       :: Maybe PreferTimezone
+    , preferMaxAffected    :: Maybe PreferMaxAffected
     , invalidPrefs         :: [ByteString]
     }
 
@@ -64,7 +67,7 @@ data Preferences
 -- >>> let sc = S.fromList ["America/Los_Angeles"]
 --
 -- One header with comma-separated values can be used to set multiple preferences:
--- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates, count=exact, timezone=America/Los_Angeles")]
+-- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates, count=exact, timezone=America/Los_Angeles, max-affected=100")]
 -- Preferences
 --     { preferResolution = Just IgnoreDuplicates
 --     , preferRepresentation = Nothing
@@ -75,12 +78,14 @@ data Preferences
 --     , preferHandling = Nothing
 --     , preferTimezone = Just
 --         ( PreferTimezone "America/Los_Angeles" )
+--     , preferMaxAffected = Just
+--         ( PreferMaxAffected 100 )
 --     , invalidPrefs = []
 --     }
 --
 -- Multiple headers can also be used:
 --
--- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates"), ("Prefer", "count=exact"), ("Prefer", "missing=null"), ("Prefer", "handling=lenient"), ("Prefer", "invalid")]
+-- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates"), ("Prefer", "count=exact"), ("Prefer", "missing=null"), ("Prefer", "handling=lenient"), ("Prefer", "invalid"), ("Prefer", "max-affected=5999")]
 -- Preferences
 --     { preferResolution = Just IgnoreDuplicates
 --     , preferRepresentation = Nothing
@@ -90,6 +95,8 @@ data Preferences
 --     , preferMissing = Just ApplyNulls
 --     , preferHandling = Just Lenient
 --     , preferTimezone = Nothing
+--     , preferMaxAffected = Just
+--         ( PreferMaxAffected 5999 )
 --     , invalidPrefs = [ "invalid" ]
 --     }
 --
@@ -121,6 +128,7 @@ data Preferences
 --     , preferMissing = Just ApplyDefaults
 --     , preferHandling = Just Strict
 --     , preferTimezone = Nothing
+--     , preferMaxAffected = Nothing
 --     , invalidPrefs = [ "anything" ]
 --     }
 --
@@ -135,7 +143,8 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     , preferMissing        = parsePrefs [ApplyDefaults, ApplyNulls]
     , preferHandling       = parsePrefs [Strict, Lenient]
     , preferTimezone       = if isTimezonePrefAccepted then PreferTimezone <$> timezonePref else Nothing
-    , invalidPrefs         = filter checkPrefs prefs
+    , preferMaxAffected    = PreferMaxAffected <$> maxAffectedPref
+    , invalidPrefs         = filter isUnacceptable prefs
     }
   where
     mapToHeadVal :: ToHeaderValue a => [a] -> [ByteString]
@@ -151,10 +160,16 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     prefHeaders = filter ((==) HTTP.hPrefer . fst) headers
     prefs = fmap BS.strip . concatMap (BS.split ',' . snd) $ prefHeaders
 
-    timezonePref = listToMaybe $ mapMaybe (BS.stripPrefix "timezone=") prefs
+    listStripPrefix prefix prefList = listToMaybe $ mapMaybe (BS.stripPrefix prefix) prefList
+
+    timezonePref = listStripPrefix "timezone=" prefs
     isTimezonePrefAccepted = (S.member <$> timezonePref <*> pure acceptedTzNames) == Just True
 
-    checkPrefs p = p `notElem` acceptedPrefs && not isTimezonePrefAccepted
+    maxAffectedPref = listStripPrefix "max-affected=" prefs >>= readMaybe . BS.unpack
+
+    isUnacceptable p = p `notElem` acceptedPrefs &&
+                       (isNothing (BS.stripPrefix "timezone=" p) ||  not isTimezonePrefAccepted) &&
+                       isNothing (BS.stripPrefix "max-affected=" p)
 
     parsePrefs :: ToHeaderValue a => [a] -> Maybe a
     parsePrefs vals =
@@ -164,7 +179,7 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     prefMap = Map.fromList . fmap (\pref -> (toHeaderValue pref, pref))
 
 prefAppliedHeader :: Preferences -> Maybe HTTP.Header
-prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferParameters, preferCount, preferTransaction, preferMissing, preferHandling, preferTimezone } =
+prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferParameters, preferCount, preferTransaction, preferMissing, preferHandling, preferTimezone, preferMaxAffected } =
   if null prefsVals
     then Nothing
     else Just (HTTP.hPreferenceApplied, combined)
@@ -179,6 +194,7 @@ prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferPar
       , toHeaderValue <$> preferTransaction
       , toHeaderValue <$> preferHandling
       , toHeaderValue <$> preferTimezone
+      , if preferHandling == Just Strict then toHeaderValue <$> preferMaxAffected else Nothing
       ]
 
 -- |
@@ -278,3 +294,10 @@ newtype PreferTimezone = PreferTimezone ByteString
 
 instance ToHeaderValue PreferTimezone where
   toHeaderValue (PreferTimezone tz) = "timezone=" <> tz
+
+-- |
+-- Limit Affected Resources
+newtype PreferMaxAffected = PreferMaxAffected Int64
+
+instance ToHeaderValue PreferMaxAffected where
+  toHeaderValue (PreferMaxAffected n) = "max-affected=" <> show n

src/PostgREST/ApiRequest/Types.hs

@@ -96,6 +96,7 @@ data ApiRequestError
   | PutMatchingPkError
   | SingularityError Integer
   | PGRSTParseError
+  | MaxAffectedViolationError Integer
   deriving Show
 
 data QPError = QPError Text Text

src/PostgREST/Error.hs

@@ -62,35 +62,36 @@ class (JSON.ToJSON a) => PgrstError a where
     responseLBS (status err) (baseHeader : headers err) $ errorPayload err
 
 instance PgrstError ApiRequestError where
-  status AggregatesNotAllowed{}  = HTTP.status400
-  status AmbiguousRelBetween{}   = HTTP.status300
-  status AmbiguousRpc{}          = HTTP.status300
-  status MediaTypeError{}        = HTTP.status415
-  status InvalidBody{}           = HTTP.status400
-  status InvalidFilters          = HTTP.status405
-  status InvalidPreferences{}    = HTTP.status400
-  status InvalidRpcMethod{}      = HTTP.status405
-  status InvalidRange{}          = HTTP.status416
-  status NotFound                = HTTP.status404
-
-  status NoRelBetween{}          = HTTP.status400
-  status NoRpc{}                 = HTTP.status404
-  status NotEmbedded{}           = HTTP.status400
-  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
-  status LimitNoOrderError       = HTTP.status400
-  status ColumnNotFound{}        = HTTP.status400
-  status GucHeadersError         = HTTP.status500
-  status GucStatusError          = HTTP.status500
-  status OffLimitsChangesError{} = HTTP.status400
-  status PutMatchingPkError      = HTTP.status400
-  status SingularityError{}      = HTTP.status406
-  status PGRSTParseError         = HTTP.status500
+  status AggregatesNotAllowed{}      = HTTP.status400
+  status AmbiguousRelBetween{}       = HTTP.status300
+  status AmbiguousRpc{}              = HTTP.status300
+  status MediaTypeError{}            = HTTP.status415
+  status InvalidBody{}               = HTTP.status400
+  status InvalidFilters              = HTTP.status405
+  status InvalidPreferences{}        = HTTP.status400
+  status InvalidRpcMethod{}          = HTTP.status405
+  status InvalidRange{}              = HTTP.status416
+  status NotFound                    = HTTP.status404
+
+  status NoRelBetween{}              = HTTP.status400
+  status NoRpc{}                     = HTTP.status404
+  status NotEmbedded{}               = HTTP.status400
+  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
+  status LimitNoOrderError           = HTTP.status400
+  status ColumnNotFound{}            = HTTP.status400
+  status GucHeadersError             = HTTP.status500
+  status GucStatusError              = HTTP.status500
+  status OffLimitsChangesError{}     = HTTP.status400
+  status PutMatchingPkError          = HTTP.status400
+  status SingularityError{}          = HTTP.status406
+  status PGRSTParseError             = HTTP.status500
+  status MaxAffectedViolationError{} = HTTP.status400
 
   headers SingularityError{}     = [MediaType.toContentType $ MTVndSingularJSON False]
   headers _ = mempty
@@ -199,6 +200,12 @@ instance JSON.ToJSON ApiRequestError where
   toJSON AggregatesNotAllowed = toJsonPgrstError
     ApiRequestErrorCode23 "Use of aggregate functions is not allowed" Nothing Nothing
 
+  toJSON (MaxAffectedViolationError n) = toJsonPgrstError
+    ApiRequestErrorCode24
+    "Query result exceeds max-affected preference constraint"
+    (Just $ JSON.String $ T.unwords ["The query affects", show n, "rows"])
+    Nothing
+
   toJSON (NoRelBetween parent child embedHint schema allRels) = toJsonPgrstError
     SchemaCacheErrorCode00
     ("Could not find a relationship between '" <> parent <> "' and '" <> child <> "' in the schema cache")
@@ -606,6 +613,7 @@ data ErrorCode
   | ApiRequestErrorCode21
   | ApiRequestErrorCode22
   | ApiRequestErrorCode23
+  | ApiRequestErrorCode24
   -- Schema Cache errors
   | SchemaCacheErrorCode00
   | SchemaCacheErrorCode01
@@ -653,6 +661,7 @@ buildErrorCode code = "PGRST" <> case code of
   ApiRequestErrorCode21  -> "121"
   ApiRequestErrorCode22  -> "122"
   ApiRequestErrorCode23  -> "123"
+  ApiRequestErrorCode24  -> "124"
 
   SchemaCacheErrorCode00 -> "200"
   SchemaCacheErrorCode01 -> "201"

src/PostgREST/Query.hs

@@ -33,6 +33,8 @@ import qualified PostgREST.SchemaCache        as SchemaCache
 
 import PostgREST.ApiRequest              (ApiRequest (..))
 import PostgREST.ApiRequest.Preferences  (PreferCount (..),
+                                          PreferHandling (..),
+                                          PreferMaxAffected (..),
                                           PreferTimezone (..),
                                           PreferTransaction (..),
                                           Preferences (..),
@@ -114,9 +116,10 @@ createQuery mrPlan@MutateReadPlan{mrMedia} apiReq conf = do
   pure resultSet
 
 updateQuery :: MutateReadPlan -> ApiRequest -> AppConfig -> DbHandler ResultSet
-updateQuery mrPlan@MutateReadPlan{mrMedia} apiReq@ApiRequest{..} conf = do
+updateQuery mrPlan@MutateReadPlan{mrMedia} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} conf = do
   resultSet <- writeQuery mrPlan apiReq conf
   failNotSingular mrMedia resultSet
+  failExceedsMaxAffectedPref (preferMaxAffected,preferHandling) resultSet
   failsChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) resultSet
   optionalRollback conf apiReq
   pure resultSet
@@ -141,9 +144,10 @@ failPut RSStandard{rsQueryTotal=queryTotal} =
     throwError $ Error.ApiRequestError ApiRequestTypes.PutMatchingPkError
 
 deleteQuery :: MutateReadPlan -> ApiRequest -> AppConfig -> DbHandler ResultSet
-deleteQuery mrPlan@MutateReadPlan{mrMedia} apiReq@ApiRequest{..} conf = do
+deleteQuery mrPlan@MutateReadPlan{mrMedia} apiReq@ApiRequest{iPreferences=Preferences{..}, ..} conf = do
   resultSet <- writeQuery mrPlan apiReq conf
   failNotSingular mrMedia resultSet
+  failExceedsMaxAffectedPref (preferMaxAffected,preferHandling) resultSet
   failsChangesOffLimits (RangeQuery.rangeLimit iTopLevelRange) resultSet
   optionalRollback conf apiReq
   pure resultSet
@@ -165,6 +169,7 @@ invokeQuery rout CallReadPlan{..} apiReq@ApiRequest{iPreferences=Preferences{..}
 
   optionalRollback conf apiReq
   failNotSingular crMedia resultSet
+  failExceedsMaxAffectedPref (preferMaxAffected,preferHandling) resultSet
   pure resultSet
 
 openApiQuery :: SchemaCache -> PgVersion -> AppConfig -> Schema -> DbHandler (Maybe (TablesMap, RoutineMap, Maybe Text))
@@ -213,6 +218,13 @@ failNotSingular mediaType RSStandard{rsQueryTotal=queryTotal} =
     lift SQL.condemn
     throwError $ Error.ApiRequestError . ApiRequestTypes.SingularityError $ toInteger queryTotal
 
+failExceedsMaxAffectedPref :: (Maybe PreferMaxAffected, Maybe PreferHandling) -> ResultSet -> DbHandler ()
+failExceedsMaxAffectedPref (Nothing,_) _ = pure ()
+failExceedsMaxAffectedPref _ RSPlan{} = pure ()
+failExceedsMaxAffectedPref (Just (PreferMaxAffected n), handling) RSStandard{rsQueryTotal=queryTotal} = when ((queryTotal > n) && (handling == Just Strict)) $ do
+  lift SQL.condemn
+  throwError $ Error.ApiRequestError . ApiRequestTypes.MaxAffectedViolationError $ toInteger queryTotal
+
 failsChangesOffLimits :: Maybe Integer -> ResultSet -> DbHandler ()
 failsChangesOffLimits _ RSPlan{} = pure ()
 failsChangesOffLimits Nothing _  = pure ()

src/PostgREST/Response.hs

@@ -74,7 +74,7 @@ readResponse WrappedReadPlan{wrMedia} headersOnly identifier ctxApiRequest@ApiRe
     RSStandard{..} -> do
       let
         (status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal rsTableTotal
-        prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing Nothing preferCount preferTransaction Nothing preferHandling preferTimezone []
+        prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing []
         headers =
           [ contentRange
           , ( "Content-Location"
@@ -105,7 +105,7 @@ createResponse QualifiedIdentifier{..} MutateReadPlan{mrMutatePlan, mrMedia} ctx
       pkCols = case mrMutatePlan of { Insert{insPkCols} -> insPkCols; _ -> mempty;}
       prefHeader = prefAppliedHeader $
         Preferences (if null pkCols && isNothing (qsOnConflict iQueryParams) then Nothing else preferResolution)
-        preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone []
+        preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone Nothing []
       headers =
         catMaybes
           [ if null rsLocation then
@@ -146,7 +146,7 @@ updateResponse MutateReadPlan{mrMedia} ctxApiRequest@ApiRequest{iPreferences=Pre
       contentRangeHeader =
         Just . RangeQuery.contentRangeH 0 (rsQueryTotal - 1) $
           if shouldCount preferCount then Just rsQueryTotal else Nothing
-      prefHeader = prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone []
+      prefHeader = prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone preferMaxAffected []
       headers = catMaybes [contentRangeHeader, prefHeader]
 
     let (status, headers', body) =
@@ -166,7 +166,7 @@ singleUpsertResponse :: MutateReadPlan -> ApiRequest -> ResultSet -> Either Erro
 singleUpsertResponse MutateReadPlan{mrMedia} ctxApiRequest@ApiRequest{iPreferences=Preferences{..}} resultSet = case resultSet of
   RSStandard {..} -> do
     let
-      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone []
+      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing []
       cTHeader = contentTypeHeaders mrMedia ctxApiRequest
 
     let isInsertIfGTZero i = if i > 0 then HTTP.status201 else HTTP.status200
@@ -190,7 +190,7 @@ deleteResponse MutateReadPlan{mrMedia} ctxApiRequest@ApiRequest{iPreferences=Pre
       contentRangeHeader =
         RangeQuery.contentRangeH 1 0 $
           if shouldCount preferCount then Just rsQueryTotal else Nothing
-      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone []
+      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone preferMaxAffected []
       headers = contentRangeHeader : prefHeader
 
     let (status, headers', body) =
@@ -243,7 +243,7 @@ invokeResponse CallReadPlan{crMedia} invMethod proc ctxApiRequest@ApiRequest{iPr
         then Error.errorPayload $ Error.ApiRequestError $ ApiRequestTypes.InvalidRange
           $ ApiRequestTypes.OutOfBounds (show $ RangeQuery.rangeOffset iTopLevelRange) (maybe "0" show rsTableTotal)
         else LBS.fromStrict rsBody
-      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing preferParameters preferCount preferTransaction Nothing preferHandling preferTimezone []
+      prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing preferParameters preferCount preferTransaction Nothing preferHandling preferTimezone preferMaxAffected []
       headers = contentRange : prefHeader
 
     let (status', headers', body) =