add: per request statement timeout using prefer header

WIP

Signed-off-by: Taimoor Zaeem <taimoorzaeem@gmail.com>

Author: taimoorzaeem

src/PostgREST/ApiRequest/Preferences.hs

@@ -6,7 +6,7 @@
 --
 -- [1] https://datatracker.ietf.org/doc/html/rfc7240
 --
-{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE RecordWildCards #-}
 module PostgREST.ApiRequest.Preferences
   ( Preferences(..)
   , PreferCount(..)
@@ -17,6 +17,7 @@ module PostgREST.ApiRequest.Preferences
   , PreferTransaction(..)
   , PreferTimezone(..)
   , PreferMaxAffected(..)
+  , PreferTimeout(..)
   , fromHeaders
   , shouldCount
   , shouldExplainCount
@@ -43,6 +44,7 @@ import Protolude
 -- >>> deriving instance Show PreferHandling
 -- >>> deriving instance Show PreferTimezone
 -- >>> deriving instance Show PreferMaxAffected
+-- >>> deriving instance Show PreferTimeout
 -- >>> deriving instance Show Preferences
 
 -- | Preferences recognized by the application.
@@ -56,6 +58,7 @@ data Preferences
     , preferHandling       :: Maybe PreferHandling
     , preferTimezone       :: Maybe PreferTimezone
     , preferMaxAffected    :: Maybe PreferMaxAffected
+    , preferTimeout        :: Maybe PreferTimeout
     , invalidPrefs         :: [ByteString]
     }
 
@@ -77,6 +80,7 @@ data Preferences
 --         ( PreferTimezone "America/Los_Angeles" )
 --     , preferMaxAffected = Just
 --         ( PreferMaxAffected 100 )
+--     , preferTimeout = Nothing
 --     , invalidPrefs = []
 --     }
 --
@@ -93,6 +97,7 @@ data Preferences
 --     , preferTimezone = Nothing
 --     , preferMaxAffected = Just
 --         ( PreferMaxAffected 5999 )
+--     , preferTimeout = Nothing
 --     , invalidPrefs = [ "invalid" ]
 --     }
 --
@@ -124,6 +129,7 @@ data Preferences
 --     , preferHandling = Just Strict
 --     , preferTimezone = Nothing
 --     , preferMaxAffected = Nothing
+--     , preferTimeout = Nothing
 --     , invalidPrefs = [ "anything" ]
 --     }
 --
@@ -138,7 +144,8 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     , preferHandling       = parsePrefs [Strict, Lenient]
     , preferTimezone       = if isTimezonePrefAccepted then PreferTimezone <$> timezonePref else Nothing
     , preferMaxAffected    = PreferMaxAffected <$> maxAffectedPref
-    , invalidPrefs         = filter isUnacceptable prefs
+    , preferTimeout        = PreferTimeout <$> timeoutPref
+    , invalidPrefs         = filter (not . isPrefValid) prefs
     }
   where
     mapToHeadVal :: ToHeaderValue a => [a] -> [ByteString]
@@ -159,10 +166,13 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     isTimezonePrefAccepted = ((S.member . decodeUtf8 <$> timezonePref) <*> pure acceptedTzNames) == Just True
 
     maxAffectedPref = listStripPrefix "max-affected=" prefs >>= readMaybe . BS.unpack
+    timeoutPref     = listStripPrefix "timeout=" prefs >>= readMaybe . BS.unpack
 
-    isUnacceptable p = p `notElem` acceptedPrefs &&
-                       (isNothing (BS.stripPrefix "timezone=" p) ||  not isTimezonePrefAccepted) &&
-                       isNothing (BS.stripPrefix "max-affected=" p)
+    isPrefValid p =
+      p `elem` acceptedPrefs ||
+      (isJust (BS.stripPrefix "timezone=" p) && isTimezonePrefAccepted) ||
+      isJust (BS.stripPrefix "max-affected=" p) ||
+      isJust (BS.stripPrefix "timeout=" p)
 
     parsePrefs :: ToHeaderValue a => [a] -> Maybe a
     parsePrefs vals =
@@ -171,8 +181,9 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     prefMap :: ToHeaderValue a => [a] -> Map.Map ByteString a
     prefMap = Map.fromList . fmap (\pref -> (toHeaderValue pref, pref))
 
+
 prefAppliedHeader :: Preferences -> Maybe HTTP.Header
-prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferCount, preferTransaction, preferMissing, preferHandling, preferTimezone, preferMaxAffected } =
+prefAppliedHeader Preferences{..} =
   if null prefsVals
     then Nothing
     else Just (HTTP.hPreferenceApplied, combined)
@@ -187,6 +198,7 @@ prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferCou
       , toHeaderValue <$> preferHandling
       , toHeaderValue <$> preferTimezone
       , if preferHandling == Just Strict then toHeaderValue <$> preferMaxAffected else Nothing
+      , if preferHandling == Just Strict then toHeaderValue <$> preferTimeout else Nothing
       ]
 
 -- |
@@ -289,3 +301,10 @@ newtype PreferMaxAffected = PreferMaxAffected Int64
 
 instance ToHeaderValue PreferMaxAffected where
   toHeaderValue (PreferMaxAffected n) = "max-affected=" <> show n
+
+-- |
+-- Statement Timeout per request
+newtype PreferTimeout = PreferTimeout Int
+
+instance ToHeaderValue PreferTimeout where
+  toHeaderValue (PreferTimeout n) = "timeout=" <> show n

src/PostgREST/Query/PreQuery.hs

@@ -17,7 +17,9 @@ import qualified Hasql.DynamicStatements.Snippet as SQL hiding (sql)
 
 
 import PostgREST.ApiRequest              (ApiRequest (..))
-import PostgREST.ApiRequest.Preferences  (PreferTimezone (..),
+import PostgREST.ApiRequest.Preferences  (PreferHandling (..),
+                                          PreferTimeout (..),
+                                          PreferTimezone (..),
                                           Preferences (..))
 import PostgREST.Auth.Types              (AuthResult (..))
 import PostgREST.Config                  (AppConfig (..))
@@ -35,11 +37,11 @@ import Protolude hiding (Handler)
 
 -- sets transaction variables
 txVarQuery :: DbActionPlan -> AppConfig -> AuthResult -> ApiRequest -> SQL.Snippet
-txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
+txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{iPreferences=Preferences{..}, ..} =
     -- To ensure `GRANT SET ON PARAMETER <superuser_setting> TO authenticator` works, the role settings must be set before the impersonated role.
     -- Otherwise the GRANT SET would have to be applied to the impersonated role. See https://github.com/PostgREST/postgrest/issues/3045
     "select " <> intercalateSnippet ", " (
-      searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql
+      searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ timeoutSql ++ funcSettingsSql ++ appSettingsSql
     )
   where
     methodSql = setConfigWithConstantName ("request.method", iMethod)
@@ -50,7 +52,11 @@ txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
     roleSql = [setConfigWithConstantName ("role", authRole)]
     roleSettingsSql = setConfigWithDynamicName <$> HM.toList (fromMaybe mempty $ HM.lookup authRole configRoleSettings)
     appSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> configAppSettings
-    timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) $ preferTimezone iPreferences
+    timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) preferTimezone
+    timeoutSql =
+      case (preferTimeout, preferHandling) of -- only applied on when handling = strict
+        (Just (PreferTimeout t), Just Strict) -> [setConfigWithConstantName ("statement_timeout", show t <> "s")]
+        _                                     -> mempty
     funcSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> funcSettings
     searchPathSql =
       let schemas = escapeIdentList (iSchema : configDbExtraSearchPath) in

src/PostgREST/Response.hs

@@ -299,4 +299,4 @@ responsePreferences plan ApiRequest{iPreferences=Preferences{..}, iQueryParams=Q
       CallReadPlan{}                            -> preferMaxAffected
       _                                         -> Nothing
 
-    in Preferences preferResolution' preferRepresentation' preferCount preferTransaction preferMissing' preferHandling preferTimezone preferMaxAffected' []
+    in Preferences preferResolution' preferRepresentation' preferCount preferTransaction preferMissing' preferHandling preferTimezone preferMaxAffected' preferTimeout []

test/spec/Feature/Query/PreferencesSpec.hs

@@ -236,3 +236,15 @@ spec =
           `shouldRespondWith`
           [json| {"code":"PGRST128","details":null,"hint":null,"message":"Function must return SETOF or TABLE when max-affected preference is used with handling=strict"} |]
           { matchStatus = 400 }
+
+    -- Prefer: Timeout
+    context "test Prefer: timeout with handling=strict" $
+      it "should fail when timeout is more than the query time" $
+        request methodGet "/rpc/sleep?seconds=5"
+          [("Prefer", "timeout=4, handling=strict")]
+          ""
+          `shouldRespondWith`
+          [json| {"code":"57014","details":null,"hint":null,"message":"canceling statement due to statement timeout"} |]
+          { matchStatus = 500
+          , matchHeaders = [ matchContentTypeJson ] -- Should this have Preference-Applied?
+          }