add: per request statement timeout using prefer header

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

Author: taimoorzaeem

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file. From versio
 ### Added
 
 - Log error when `db-schemas` config contains schema `pg_catalog` or `information_schema` by @taimoorzaeem in #4359
+- Add `Prefer: timeout` header for per request statement timeout by @taimoorzaeem in #4381
 
 ## [14.2] - 2025-12-18
 

docs/references/api/preferences.rst

@@ -15,6 +15,7 @@ The following preferences are supported.
 - ``Prefer: missing``. See :ref:`prefer_missing`.
 - ``Prefer: max-affected``, See :ref:`prefer_max_affected`.
 - ``Prefer: tx``. See :ref:`prefer_tx`.
+- ``Prefer: timeout``. See :ref:`prefer_timeout`.
 
 .. _prefer_handling:
 
@@ -296,3 +297,38 @@ With :ref:`RPC <functions>`, the preference is honored completely on the basis o
 .. note::
 
   It is important for functions to return ``SETOF`` or ``TABLE`` when called with ``max-affected`` preference. A violation of this would cause a :ref:`PGRST128 <pgrst128>` error.
+
+
+.. _prefer_timeout:
+
+Timeout
+=======
+
+You can set `statement_timeout <https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-STATEMENT-TIMEOUT>`_ for the request using this preference. This works in combination with ``handling=strict`` preference in the same header.
+
+Currently, the header only accepts integer value indicating the ``seconds`` that are set as timeout value. To demonstrate, see the following example:
+
+.. code-block:: postgres
+
+  CREATE FUNCTION test.sleep(seconds)
+  RETURNS VOID AS $$
+    SELECT pg_sleep(seconds);
+  $$ LANGUAGE SQL;
+
+.. code-block:: bash
+
+  curl -i "http://localhost:3000/rpc/sleep?seconds=5" \
+    -H "Prefer: handling=strict, timeout=2"
+
+.. code-block:: http
+
+  HTTP/1.1 500 Internal Server Error
+
+.. code-block:: json
+
+  {
+      "code": "57014",
+      "details": null,
+      "hint": null,
+      "message": "canceling statement due to statement timeout"
+  }

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,63 @@ 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 }
+
+    context "Prefer: timeout and handling=strict" $ do
+      it "should fail when timeout is less than the query time" $
+        request methodGet "/rpc/sleep?seconds=4"
+          [("Prefer", "handling=strict, timeout=3")]
+          ""
+          `shouldRespondWith`
+          [json| {"code":"57014","details":null,"hint":null,"message":"canceling statement due to statement timeout"} |]
+          { matchStatus = 500
+          , matchHeaders = [ matchContentTypeJson ]
+          }
+
+      it "should fail when timeout is equal to the query time" $
+        request methodGet "/rpc/sleep?seconds=3"
+          [("Prefer", "handling=strict, timeout=3")]
+          ""
+          `shouldRespondWith`
+          [json| {"code":"57014","details":null,"hint":null,"message":"canceling statement due to statement timeout"} |]
+          { matchStatus = 500
+          , matchHeaders = [ matchContentTypeJson ]
+          }
+
+      it "should succeed when timeout is more than the query time" $
+        request methodGet "/rpc/sleep?seconds=3"
+          [("Prefer", "handling=strict, timeout=4")]
+          ""
+          `shouldRespondWith`
+          ""
+          { matchStatus = 204
+          , matchHeaders = ["Preference-Applied" <:> "handling=strict, timeout=4"]
+          }
+
+    context "Prefer: timeout and handling=lenient" $
+      it "statement timeout is not applied when handling=lenient, so it should succeed in all cases" $ do
+        request methodGet "/rpc/sleep?seconds=4"
+          [("Prefer", "handling=lenient, timeout=3")]
+          ""
+          `shouldRespondWith`
+          ""
+          { matchStatus = 204
+          , matchHeaders = ["Preference-Applied" <:> "handling=lenient"]
+          }
+
+        request methodGet "/rpc/sleep?seconds=3"
+          [("Prefer", "handling=lenient, timeout=3")]
+          ""
+          `shouldRespondWith`
+          ""
+          { matchStatus = 204
+          , matchHeaders = ["Preference-Applied" <:> "handling=lenient"]
+          }
+
+        request methodGet "/rpc/sleep?seconds=3"
+          [("Prefer", "handling=lenient, timeout=4")]
+          ""
+          `shouldRespondWith`
+          ""
+          { matchStatus = 204
+          , matchHeaders = ["Preference-Applied" <:> "handling=lenient"]
+          }