add: per request statement timeout using prefer header

Adds a feature to set `statement_timeout` using `Prefer: timeout`
header. This also introduces a `PGRST129` error which is returned
when the timeout preferred exceeds the per-role or global timeout.

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

Author: taimoorzaeem

CHANGELOG.md

@@ -12,6 +12,7 @@ All notable changes to this project will be documented in this file. From versio
 - Log host, port and pg version of listener database connection by @mkleczek in #4617 #4618
 - Optimize requests with `Prefer: count=exact` that do not use ranges or `db-max-rows` by @laurenceisla in #3957
   + Removed unnecessary double count when building the `Content-Range`.
+- Add `Prefer: timeout` header for per-request `statement_timeout` by @taimoorzaeem in #4381
 
 ### Changed
 

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,62 @@ 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.
+
+The header only accepts integer value indicating the ``seconds`` that are set as timeout value. ``0`` and negative values are ignored and not applied. 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"
+  }
+
+It is important to note the timeout value cannot exceed the ``statement_timeout`` set :ref:`per-role <impersonated_settings>` or per-database. The role's timeout setting takes precedence over the database level timeout. This restriction prevents misuse of this feature. PostgREST returns a :ref:`PGRST129 <pgrst129>` error in this case.
+
+.. code-block:: postgres
+
+  ALTER ROLE postgrest_test_anonymous SET statement_timeout = '3s';
+
+.. code-block:: bash
+
+  curl -i "http://localhost:3000/rpc/sleep?seconds=4" \
+    -H "Prefer: handling=strict, timeout=5"
+
+.. code-block:: http
+
+  HTTP/1.1 400 Bad Request
+  Content-Type: application/json; charset=utf-8
+
+.. code-block:: json
+
+  {
+      "code": "PGRST129",
+      "message": "Timeout preference exceeded maximum allowed",
+      "details": "The maximum timeout allowed is 3s",
+      "hint": "Reduce the timeout"
+  }

docs/references/errors.rst

@@ -271,6 +271,10 @@ Related to the HTTP request elements.
 |               |             | See :ref:`prefer_max_affected`.                             |
 | PGRST128      |             |                                                             |
 +---------------+-------------+-------------------------------------------------------------+
+| .. _pgrst129: | 400         | ``timeout`` preference exceeds ``statement_timeout`` value  |
+|               |             | of role. See :ref:`prefer_timeout`.                         |
+| PGRST129      |             |                                                             |
++---------------+-------------+-------------------------------------------------------------+
 
 
 .. _pgrst2**:

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,12 +80,13 @@ data Preferences
 --         ( PreferTimezone "America/Los_Angeles" )
 --     , preferMaxAffected = Just
 --         ( PreferMaxAffected 100 )
+--     , preferTimeout = Nothing
 --     , 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"), ("Prefer", "max-affected=5999")]
+-- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates"), ("Prefer", "count=exact"), ("Prefer", "missing=null"), ("Prefer", "handling=lenient"), ("Prefer", "invalid"), ("Prefer", "max-affected=5999"), ("Prefer", "timeout=10")]
 -- Preferences
 --     { preferResolution = Just IgnoreDuplicates
 --     , preferRepresentation = Nothing
@@ -93,6 +97,8 @@ data Preferences
 --     , preferTimezone = Nothing
 --     , preferMaxAffected = Just
 --         ( PreferMaxAffected 5999 )
+--     , preferTimeout = Just
+--         ( PreferTimeout 10 )
 --     , invalidPrefs = [ "invalid" ]
 --     }
 --
@@ -124,6 +130,7 @@ data Preferences
 --     , preferHandling = Just Strict
 --     , preferTimezone = Nothing
 --     , preferMaxAffected = Nothing
+--     , preferTimeout = Nothing
 --     , invalidPrefs = [ "anything" ]
 --     }
 --
@@ -138,7 +145,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 +167,17 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     isTimezonePrefAccepted = ((S.member . decodeUtf8 <$> timezonePref) <*> pure acceptedTzNames) == Just True
 
     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)
+    timeoutPref = listStripPrefix "timeout=" prefs >>= readOnlyPositiveInt . readMaybe . BS.unpack
+      where
+        -- 0 and -ve values for timeout are meaningless, we handle them leniently and ignore them
+        readOnlyPositiveInt (Just i) | i > 0 = Just i
+        readOnlyPositiveInt _ = Nothing
+
+    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 =
@@ -172,7 +187,7 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
     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 +202,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 +305,10 @@ newtype PreferMaxAffected = PreferMaxAffected Int64
 
 instance ToHeaderValue PreferMaxAffected where
   toHeaderValue (PreferMaxAffected n) = "max-affected=" <> show n
+
+-- |
+-- Statement Timeout per request
+newtype PreferTimeout = PreferTimeout Int64
+
+instance ToHeaderValue PreferTimeout where
+  toHeaderValue (PreferTimeout n) = "timeout=" <> show n

src/PostgREST/App.hs

@@ -169,7 +169,7 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache authResult@AuthRe
       prefs = ApiRequest.userPreferences conf req timezones
 
   (parseTime, apiReq@ApiRequest{..}) <- withTiming $ liftEither . mapLeft Error.ApiRequestError $ ApiRequest.userApiRequest conf prefs req body
-  (planTime, plan)                   <- withTiming $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
+  (planTime, plan)                   <- withTiming $ liftEither $ Plan.actionPlan iAction conf apiReq authResult sCache
 
   let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest
       tx = MainTx.mainTx mainQ conf authResult apiReq plan sCache

src/PostgREST/AppState.hs

@@ -63,7 +63,8 @@ import PostgREST.Config                  (AppConfig (..),
                                           readAppConfig)
 import PostgREST.Config.Database         (queryDbSettings,
                                           queryPgVersion,
-                                          queryRoleSettings)
+                                          queryRoleSettings,
+                                          queryRoleTimeoutSettings)
 import PostgREST.Config.PgVersion        (PgVersion (..),
                                           minimumPgVersion)
 import PostgREST.SchemaCache             (SchemaCache (..),
@@ -461,7 +462,17 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
         Right x -> pure x
     else
       pure mempty
-  readAppConfig dbSettings (configFilePath conf) (Just $ configDbUri conf) roleSettings roleIsolationLvl >>= \case
+  roleTimeoutSettings <-
+    if configDbConfig conf then do
+      rSettings <- usePool appState (queryRoleTimeoutSettings $ configDbPreparedStatements conf)
+      case rSettings of
+        Left e -> do
+          observer $ QueryRoleSettingsErrorObs e
+          pure mempty
+        Right x -> pure x
+    else
+      pure mempty
+  readAppConfig dbSettings (configFilePath conf) (Just $ configDbUri conf) roleSettings roleTimeoutSettings roleIsolationLvl >>= \case
     Left err   ->
       if startingUp then
         panic err -- die on invalid config if the program is starting up

src/PostgREST/CLI.hs

@@ -30,7 +30,7 @@ import Protolude
 main :: CLI -> IO ()
 main CLI{cliCommand, cliPath} = do
   conf <-
-    either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
+    either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty mempty
   case cliCommand of
     Client adminCmd -> runClientCommand conf adminCmd
     Run runCmd      -> runAppCommand conf runCmd

src/PostgREST/Config.hs

@@ -57,7 +57,8 @@ import System.Environment      (getEnvironment)
 import System.Posix.Types      (FileMode)
 
 import PostgREST.Config.Database         (RoleIsolationLvl,
-                                          RoleSettings)
+                                          RoleSettings,
+                                          RoleTimeoutSettings)
 import PostgREST.Config.JSPath           (FilterExp (..), JSPath,
                                           JSPathExp (..), dumpJSPath,
                                           pRoleClaimKey)
@@ -117,6 +118,10 @@ data AppConfig = AppConfig
   , configAdminServerHost          :: Text
   , configAdminServerPort          :: Maybe Int
   , configRoleSettings             :: RoleSettings
+  -- Cached statement timeout settings converted to number of seconds. They
+  -- are never applied, only used to check max allowed timeout for a role
+  -- when "Prefer: timeout=" header is used.
+  , configRoleTimeoutSettings      :: RoleTimeoutSettings
   , configRoleIsoLvl               :: RoleIsolationLvl
   , configInternalSCQuerySleep     :: Maybe Int32
   , configInternalSCLoadSleep      :: Maybe Int32
@@ -227,13 +232,13 @@ instance JustIfMaybe a (Maybe a) where
 
 -- | Reads and parses the config and overrides its parameters from env vars,
 -- files or db settings.
-readAppConfig :: [(Text, Text)] -> Maybe FilePath -> Maybe Text -> RoleSettings -> RoleIsolationLvl -> IO (Either Text AppConfig)
-readAppConfig dbSettings optPath prevDbUri roleSettings roleIsolationLvl = do
+readAppConfig :: [(Text, Text)] -> Maybe FilePath -> Maybe Text -> RoleSettings -> RoleTimeoutSettings -> RoleIsolationLvl -> IO (Either Text AppConfig)
+readAppConfig dbSettings optPath prevDbUri roleSettings roleTimeoutSettings roleIsolationLvl = do
   env <- readPGRSTEnvironment
   -- if no filename provided, start with an empty map to read config from environment
   conf <- maybe (return $ Right M.empty) loadConfig optPath
 
-  case C.runParser (parser optPath env dbSettings roleSettings roleIsolationLvl) =<< mapLeft show conf of
+  case C.runParser (parser optPath env dbSettings roleSettings roleTimeoutSettings roleIsolationLvl) =<< mapLeft show conf of
     Left err ->
       return . Left $ "Error in config " <> err
     Right parsedConfig ->
@@ -250,8 +255,8 @@ readAppConfig dbSettings optPath prevDbUri roleSettings roleIsolationLvl = do
       readSecretFile =<<
       readDbUriFile prevDbUri parsedConfig
 
-parser :: Maybe FilePath -> Environment -> [(Text, Text)] -> RoleSettings -> RoleIsolationLvl -> C.Parser C.Config AppConfig
-parser optPath env dbSettings roleSettings roleIsolationLvl =
+parser :: Maybe FilePath -> Environment -> [(Text, Text)] -> RoleSettings -> RoleTimeoutSettings -> RoleIsolationLvl -> C.Parser C.Config AppConfig
+parser optPath env dbSettings roleSettings roleTimeoutSettings roleIsolationLvl =
   AppConfig
     <$> parseAppSettings "app.settings"
     <*> (fromMaybe False <$> optBool "db-aggregates-enabled")
@@ -305,6 +310,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
                                             (optString "server-host"))
     <*> parseAdminServerPort "admin-server-port"
     <*> pure roleSettings
+    <*> pure roleTimeoutSettings
     <*> pure roleIsolationLvl
     <*> optInt "internal-schema-cache-query-sleep"
     <*> optInt "internal-schema-cache-load-sleep"