Skip to content

Commit

Permalink
feat: add metrics=timings to prefer header
Browse files Browse the repository at this point in the history
  • Loading branch information
taimoorzaeem committed May 12, 2024
1 parent b8fe512 commit 953a808
Show file tree
Hide file tree
Showing 8 changed files with 82 additions and 65 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
+ Shows the correct JSON format in the `hints` field
- #3340, Log when the LISTEN channel gets a notification - @steve-chavez
- #3184, Log full pg version to stderr on connection - @steve-chavez
- #3410, Add `metrics=timings` to Prefer header - @taimoorzaeem

### Fixed

Expand Down
24 changes: 21 additions & 3 deletions src/PostgREST/ApiRequest/Preferences.hs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ module PostgREST.ApiRequest.Preferences
, PreferTransaction(..)
, PreferTimezone(..)
, PreferMaxAffected(..)
, PreferMetrics(..)
, fromHeaders
, shouldCount
, prefAppliedHeader
Expand All @@ -44,6 +45,7 @@ import Protolude
-- >>> deriving instance Show PreferHandling
-- >>> deriving instance Show PreferTimezone
-- >>> deriving instance Show PreferMaxAffected
-- >>> deriving instance Show PreferMetrics
-- >>> deriving instance Show Preferences

-- | Preferences recognized by the application.
Expand All @@ -58,6 +60,7 @@ data Preferences
, preferHandling :: Maybe PreferHandling
, preferTimezone :: Maybe PreferTimezone
, preferMaxAffected :: Maybe PreferMaxAffected
, preferMetrics :: Maybe PreferMetrics
, invalidPrefs :: [ByteString]
}

Expand All @@ -67,7 +70,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, max-affected=100")]
-- >>> pPrint $ fromHeaders True sc [("Prefer", "resolution=ignore-duplicates, count=exact, timezone=America/Los_Angeles, max-affected=100, metrics=timings")]
-- Preferences
-- { preferResolution = Just IgnoreDuplicates
-- , preferRepresentation = Nothing
Expand All @@ -80,6 +83,7 @@ data Preferences
-- ( PreferTimezone "America/Los_Angeles" )
-- , preferMaxAffected = Just
-- ( PreferMaxAffected 100 )
-- , preferMetrics = Just Timings
-- , invalidPrefs = []
-- }
--
Expand All @@ -97,6 +101,7 @@ data Preferences
-- , preferTimezone = Nothing
-- , preferMaxAffected = Just
-- ( PreferMaxAffected 5999 )
-- , preferMetrics = Nothing
-- , invalidPrefs = [ "invalid" ]
-- }
--
Expand Down Expand Up @@ -129,6 +134,7 @@ data Preferences
-- , preferHandling = Just Strict
-- , preferTimezone = Nothing
-- , preferMaxAffected = Nothing
-- , preferMetrics = Nothing
-- , invalidPrefs = [ "anything" ]
-- }
--
Expand All @@ -144,6 +150,7 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
, preferHandling = parsePrefs [Strict, Lenient]
, preferTimezone = if isTimezonePrefAccepted then PreferTimezone <$> timezonePref else Nothing
, preferMaxAffected = PreferMaxAffected <$> maxAffectedPref
, preferMetrics = parsePrefs [Timings]
, invalidPrefs = filter isUnacceptable prefs
}
where
Expand All @@ -155,7 +162,8 @@ fromHeaders allowTxDbOverride acceptedTzNames headers =
mapToHeadVal [ExactCount, PlannedCount, EstimatedCount] ++
mapToHeadVal [Commit, Rollback] ++
mapToHeadVal [ApplyDefaults, ApplyNulls] ++
mapToHeadVal [Strict, Lenient]
mapToHeadVal [Strict, Lenient] ++
mapToHeadVal [Timings]

prefHeaders = filter ((==) HTTP.hPrefer . fst) headers
prefs = fmap BS.strip . concatMap (BS.split ',' . snd) $ prefHeaders
Expand All @@ -179,7 +187,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, preferMaxAffected } =
prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferParameters, preferCount, preferTransaction, preferMissing, preferHandling, preferTimezone, preferMaxAffected, preferMetrics } =
if null prefsVals
then Nothing
else Just (HTTP.hPreferenceApplied, combined)
Expand All @@ -195,6 +203,7 @@ prefAppliedHeader Preferences {preferResolution, preferRepresentation, preferPar
, toHeaderValue <$> preferHandling
, toHeaderValue <$> preferTimezone
, if preferHandling == Just Strict then toHeaderValue <$> preferMaxAffected else Nothing
, toHeaderValue <$> preferMetrics
]

-- |
Expand Down Expand Up @@ -302,3 +311,12 @@ newtype PreferMaxAffected = PreferMaxAffected Int64

instance ToHeaderValue PreferMaxAffected where
toHeaderValue (PreferMaxAffected n) = "max-affected=" <> show n

-- |
-- Show Performance Metrics
data PreferMetrics
= Timings -- show server timings for the request
deriving Eq

instance ToHeaderValue PreferMetrics where
toHeaderValue Timings = "metrics=timings"
61 changes: 33 additions & 28 deletions src/PostgREST/App.hs
Original file line number Diff line number Diff line change
Expand Up @@ -39,17 +39,20 @@ import qualified PostgREST.Query as Query
import qualified PostgREST.Response as Response
import qualified PostgREST.Unix as Unix (installSignalHandlers)

import PostgREST.ApiRequest (ApiRequest (..))
import PostgREST.AppState (AppState)
import PostgREST.Auth (AuthResult (..))
import PostgREST.Config (AppConfig (..), LogLevel (..))
import PostgREST.Config.PgVersion (PgVersion (..))
import PostgREST.Error (Error)
import PostgREST.Observation (Observation (..))
import PostgREST.Response.Performance (ServerTiming (..),
serverTimingHeader)
import PostgREST.SchemaCache (SchemaCache (..))
import PostgREST.Version (docsVersion, prettyVersion)
import PostgREST.ApiRequest (ApiRequest (..))
import PostgREST.ApiRequest.Preferences (PreferMetrics (..),
Preferences (..))
import PostgREST.AppState (AppState)
import PostgREST.Auth (AuthResult (..))
import PostgREST.Config (AppConfig (..),
LogLevel (..))
import PostgREST.Config.PgVersion (PgVersion (..))
import PostgREST.Error (Error)
import PostgREST.Observation (Observation (..))
import PostgREST.Response.Performance (ServerTiming (..),
serverTimingHeader)
import PostgREST.SchemaCache (SchemaCache (..))
import PostgREST.Version (docsVersion, prettyVersion)

import qualified Data.ByteString.Char8 as BS
import qualified Data.List as L
Expand Down Expand Up @@ -141,27 +144,29 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache pgVer authResult@

body <- lift $ Wai.strictRequestBody req

let jwtTime = if configServerTimingEnabled then Auth.getJwtDur req else Nothing
-- the preference metrics=timings cant be used before it is parsed, hence
-- parseTime will be calculated for all requests
(parseTime, apiReq@ApiRequest{..}) <- withTiming True $ liftEither . mapLeft Error.ApiRequestError $ ApiRequest.userApiRequest conf req body sCache
let timingsPref = preferMetrics iPreferences == Just Timings
(planTime, plan) <- withTiming timingsPref $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
(queryTime, queryResult) <- withTiming timingsPref $ Query.runQuery appState conf authResult apiReq plan sCache pgVer (Just authRole /= configDbAnonRole)
(respTime, resp) <- withTiming timingsPref $ liftEither $ Response.actionResponse queryResult apiReq (T.decodeUtf8 prettyVersion, docsVersion) conf sCache iSchema iNegotiatedByProfile

(parseTime, apiReq@ApiRequest{..}) <- withTiming $ liftEither . mapLeft Error.ApiRequestError $ ApiRequest.userApiRequest conf req body sCache
(planTime, plan) <- withTiming $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
(queryTime, queryResult) <- withTiming $ Query.runQuery appState conf authResult apiReq plan sCache pgVer (Just authRole /= configDbAnonRole)
(respTime, resp) <- withTiming $ liftEither $ Response.actionResponse queryResult apiReq (T.decodeUtf8 prettyVersion, docsVersion) conf sCache iSchema iNegotiatedByProfile
let jwtTime = if timingsPref then Auth.getJwtDur req else Nothing

return $ toWaiResponse (ServerTiming jwtTime parseTime planTime queryTime respTime) resp
return $ toWaiResponse timingsPref (ServerTiming jwtTime parseTime planTime queryTime respTime) resp

where
toWaiResponse :: ServerTiming -> Response.PgrstResponse -> Wai.Response
toWaiResponse timing (Response.PgrstResponse st hdrs bod) = Wai.responseLBS st (hdrs ++ ([serverTimingHeader timing | configServerTimingEnabled])) bod

withTiming :: Handler IO a -> Handler IO (Maybe Double, a)
withTiming f = if configServerTimingEnabled
then do
(t, r) <- timeItT f
pure (Just t, r)
else do
r <- f
pure (Nothing, r)
toWaiResponse :: Bool -> ServerTiming -> Response.PgrstResponse -> Wai.Response
toWaiResponse includeTimings timing (Response.PgrstResponse st hdrs bod) = Wai.responseLBS st (hdrs ++ ([serverTimingHeader timing | includeTimings])) bod

withTiming :: Bool -> Handler IO a -> Handler IO (Maybe Double, a)
withTiming True f = do
(t, r) <- timeItT f
pure (Just t, r)
withTiming False f = do
r <- f
pure (Nothing, r)

traceHeaderMiddleware :: AppState -> Wai.Middleware
traceHeaderMiddleware appState app req respond = do
Expand Down
17 changes: 4 additions & 13 deletions src/PostgREST/Auth.hs
Original file line number Diff line number Diff line change
Expand Up @@ -107,25 +107,16 @@ middleware appState app req respond = do
let token = fromMaybe "" $ Wai.extractBearerAuth =<< lookup HTTP.hAuthorization (Wai.requestHeaders req)
parseJwt = runExceptT $ parseToken conf (LBS.fromStrict token) time >>= parseClaims conf

-- If DbPlanEnabled -> calculate JWT validation time
-- If JwtCacheMaxLifetime -> cache JWT validation result
req' <- case (configServerTimingEnabled conf, configJwtCacheMaxLifetime conf) of
(True, 0) -> do
-- If JwtCacheMaxLifetime -> cache JWT validation result
req' <- case configJwtCacheMaxLifetime conf of
0 -> do
(dur, authResult) <- timeItT parseJwt
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }

(True, maxLifetime) -> do
maxLifetime -> do
(dur, authResult) <- timeItT $ getJWTFromCache appState token maxLifetime parseJwt time
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult & Vault.insert jwtDurKey dur }

(False, 0) -> do
authResult <- parseJwt
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }

(False, maxLifetime) -> do
authResult <- getJWTFromCache appState token maxLifetime parseJwt time
return $ req { Wai.vault = Wai.vault req & Vault.insert authResultKey authResult }

app req' respond

-- | Used to retrieve and insert JWT to JWT Cache
Expand Down
12 changes: 6 additions & 6 deletions src/PostgREST/Response.hs
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ actionResponse (DbCrudResult WrappedReadPlan{wrMedia, wrHdrsOnly=headersOnly, cr
RSStandard{..} -> do
let
(status, contentRange) = RangeQuery.rangeStatusHeader iTopLevelRange rsQueryTotal rsTableTotal
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing []
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing preferMetrics []
headers =
[ contentRange
, ( "Content-Location"
Expand Down Expand Up @@ -99,7 +99,7 @@ actionResponse (DbCrudResult MutateReadPlan{mrMutation=MutationCreate, mrMutateP
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 Nothing []
preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone Nothing preferMetrics []
headers =
catMaybes
[ if null rsLocation then
Expand Down Expand Up @@ -139,7 +139,7 @@ actionResponse (DbCrudResult MutateReadPlan{mrMutation=MutationUpdate, mrMedia}
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 preferMaxAffected []
prefHeader = prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction preferMissing preferHandling preferTimezone preferMaxAffected preferMetrics []
headers = catMaybes [contentRangeHeader, prefHeader]

let (status, headers', body) =
Expand All @@ -158,7 +158,7 @@ actionResponse (DbCrudResult MutateReadPlan{mrMutation=MutationUpdate, mrMedia}
actionResponse (DbCrudResult MutateReadPlan{mrMutation=MutationSingleUpsert, mrMedia} resultSet) ctxApiRequest@ApiRequest{iPreferences=Preferences{..}} _ _ _ _ _ = case resultSet of
RSStandard {..} -> do
let
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing []
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone Nothing preferMetrics []
cTHeader = contentTypeHeaders mrMedia ctxApiRequest

let isInsertIfGTZero i = if i > 0 then HTTP.status201 else HTTP.status200
Expand All @@ -181,7 +181,7 @@ actionResponse (DbCrudResult MutateReadPlan{mrMutation=MutationDelete, mrMedia}
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 preferMaxAffected []
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing preferRepresentation Nothing preferCount preferTransaction Nothing preferHandling preferTimezone preferMaxAffected preferMetrics []
headers = contentRangeHeader : prefHeader

let (status, headers', body) =
Expand All @@ -206,7 +206,7 @@ actionResponse (DbCallResult CallReadPlan{crMedia, crInvMthd=invMethod, crProc=p
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 preferMaxAffected []
prefHeader = maybeToList . prefAppliedHeader $ Preferences Nothing Nothing preferParameters preferCount preferTransaction Nothing preferHandling preferTimezone preferMaxAffected preferMetrics []
headers = contentRange : prefHeader

let (status', headers', body) =
Expand Down
5 changes: 3 additions & 2 deletions test/io/test_big_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,18 @@ def test_requests_wait_for_schema_cache_reload(defaultenv):
"PGRST_DB_SCHEMAS": "apflora",
"PGRST_DB_POOL": "2",
"PGRST_DB_ANON_ROLE": "postgrest_test_anonymous",
"PGRST_SERVER_TIMING_ENABLED": "true",
}

headers = {'Prefer': 'metrics=timings'}

with run(env=env, wait_max_seconds=30) as postgrest:
# reload the schema cache
response = postgrest.session.get("/rpc/notify_pgrst")
assert response.status_code == 204

postgrest.wait_until_scache_starts_loading()

response = postgrest.session.get("/tpopmassn?select=*,tpop(*)")
response = postgrest.session.get("/tpopmassn?select=*,tpop(*)", headers=headers)
assert response.status_code == 200

plan_dur = parse_server_timings_header(response.headers["Server-Timing"])[
Expand Down
6 changes: 3 additions & 3 deletions test/io/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -1174,7 +1174,6 @@ def test_server_timing_jwt_should_decrease_on_subsequent_requests(defaultenv):

env = {
**defaultenv,
"PGRST_SERVER_TIMING_ENABLED": "true",
"PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
"PGRST_JWT_SECRET": "@/dev/stdin",
"PGRST_DB_CONFIG": "false",
Expand All @@ -1189,6 +1188,7 @@ def test_server_timing_jwt_should_decrease_on_subsequent_requests(defaultenv):
},
SECRET,
)
headers["Prefer"] = "metrics=timings"

with run(stdin=SECRET.encode(), env=env) as postgrest:
first_timings = postgrest.session.get("/authors_only", headers=headers).headers[
Expand Down Expand Up @@ -1234,13 +1234,13 @@ def test_server_timing_jwt_should_not_decrease_when_caching_disabled(defaultenv)

env = {
**defaultenv,
"PGRST_SERVER_TIMING_ENABLED": "true",
"PGRST_JWT_CACHE_MAX_LIFETIME": "0", # cache disabled
"PGRST_JWT_SECRET": "@/dev/stdin",
"PGRST_DB_CONFIG": "false",
}

headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)
headers["Prefer"] = "metrics=timings"

with run(stdin=SECRET.encode(), env=env) as postgrest:
warmup_req = postgrest.session.get("/authors_only", headers=headers)
Expand All @@ -1264,13 +1264,13 @@ def test_jwt_cache_with_no_exp_claim(defaultenv):

env = {
**defaultenv,
"PGRST_SERVER_TIMING_ENABLED": "true",
"PGRST_JWT_CACHE_MAX_LIFETIME": "86400",
"PGRST_JWT_SECRET": "@/dev/stdin",
"PGRST_DB_CONFIG": "false",
}

headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET) # no exp
headers["Prefer"] = "metrics=timings"

with run(stdin=SECRET.encode(), env=env) as postgrest:
first_timings = postgrest.session.get("/authors_only", headers=headers).headers[
Expand Down
Loading

0 comments on commit 953a808

Please sign in to comment.