refactor(test): provide means to validate metrics and observations

mkleczek · mkleczek · commit 5b7652ebf688 · 2026-03-06T06:41:42.000+01:00
DISCLAIMER:
This commit was authored entirely by a human without the assistance of LLMs.

Some helpers are provided for introspecting metrics already (used in JWT cache tests). This change provides facilities to additionally validate emited Observation events.
A new Spec module is also implemented, adding basic tests of schema cache reloading - their main goal is to excercise the new infrastructure.
diff --git a/postgrest.cabal b/postgrest.cabal
@@ -219,6 +219,7 @@ test-suite spec
                       Feature.ConcurrentSpec
                       Feature.CorsSpec
                       Feature.ExtraSearchPathSpec
+                      Feature.MetricsSpec
                       Feature.NoSuperuserSpec
                       Feature.ObservabilitySpec
                       Feature.OpenApi.DisabledOpenApiSpec
diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
@@ -18,6 +18,7 @@ module PostgREST.AppState
   , init
   , initSockets
   , initWithPool
+  , putConfig -- For tests TODO refactoring
   , putNextListenerDelay
   , putSchemaCache
   , putPgVersion
diff --git a/src/PostgREST/Observation.hs b/src/PostgREST/Observation.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE DeriveGeneric #-}
 {-|
 Module      : PostgREST.Observation
 Description : This module holds an Observation type which is the core of Observability for PostgREST.
@@ -55,6 +56,7 @@ data Observation
   | JwtCacheLookup Bool
   | JwtCacheEviction
   | WarpErrorObs Text
+  deriving (Generic)
 
 data ObsFatalError = ServerAuthError | ServerPgrstBug | ServerError42P05 | ServerError08P01
 
diff --git a/test/spec/Feature/MetricsSpec.hs b/test/spec/Feature/MetricsSpec.hs
@@ -0,0 +1,54 @@
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE MonadComprehensions #-}
+{-# LANGUAGE NamedFieldPuns      #-}
+{-# LANGUAGE TypeApplications    #-}
+
+module Feature.MetricsSpec where
+
+import           Data.List             (lookup)
+import           Network.Wai           (Application)
+import qualified PostgREST.AppState    as AppState
+import           PostgREST.Config      (AppConfig (configDbSchemas))
+import qualified PostgREST.Metrics     as Metrics
+import           PostgREST.Observation
+import           Prometheus            (getCounter, getVectorWith)
+import           Protolude
+import           SpecHelper
+import           Test.Hspec            (SpecWith, describe, it)
+import           Test.Hspec.Wai        (getState)
+
+spec :: SpecWith (SpecState, Application)
+spec = describe "Server started with metrics enabled" $ do
+  it "Should update pgrst_schema_cache_loads_total[SUCCESS]" $ do
+    SpecState{specAppState = appState, specMetrics = metrics, specObsChan} <- getState
+    let waitFor = waitForObs specObsChan
+
+    liftIO $ checkState' metrics [
+        schemaCacheLoads "SUCCESS" (+1)
+      ] $ do
+        AppState.schemaCacheLoader appState
+        waitFor (1 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@(SchemaCacheLoadedObs{}) <- pure x]
+
+  it "Should update pgrst_schema_cache_loads_total[ERROR]" $ do
+    SpecState{specAppState = appState, specMetrics = metrics, specObsChan} <- getState
+    let waitFor = waitForObs specObsChan
+
+    liftIO $ checkState' metrics [
+        schemaCacheLoads "FAIL" (+1),
+        schemaCacheLoads "SUCCESS" (+1)
+      ] $ do
+        AppState.getConfig appState >>= \prev -> do
+          AppState.putConfig appState $ prev { configDbSchemas = pure "bad_schema" }
+          AppState.schemaCacheLoader appState
+          waitFor (1 * sec) "SchemaCacheErrorObs" $ \x -> [ o | o@(SchemaCacheErrorObs{}) <- pure x]
+          AppState.putConfig appState prev
+
+        -- wait up to 2 secs so that retry can happen
+        waitFor (2 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@(SchemaCacheLoadedObs{}) <- pure x]
+
+  where
+    -- prometheus-client api to handle vectors is convoluted
+    schemaCacheLoads label = expectField @"schemaCacheLoads" $
+      fmap (maybe (0::Int) round . lookup label) . (`getVectorWith` getCounter)
+    sec = 1000000
diff --git a/test/spec/Main.hs b/test/spec/Main.hs
@@ -29,6 +29,7 @@ import qualified Feature.Auth.NoJwtSecretSpec
 import qualified Feature.ConcurrentSpec
 import qualified Feature.CorsSpec
 import qualified Feature.ExtraSearchPathSpec
+import qualified Feature.MetricsSpec
 import qualified Feature.NoSuperuserSpec
 import qualified Feature.ObservabilitySpec
 import qualified Feature.OpenApi.DisabledOpenApiSpec
@@ -68,16 +69,26 @@ import qualified Feature.Query.UpdateSpec
 import qualified Feature.Query.UpsertSpec
 import qualified Feature.RollbackSpec
 import qualified Feature.RpcPreRequestGucsSpec
+import           PostgREST.Observation                 (Observation (HasqlPoolObs))
 
 
 main :: IO ()
 main = do
+  poolChan <- newChan
+  -- make sure poolChan is not growing indefinitely
+  -- start a thread that drains the channel
+  -- this is necessary because test cases operate on
+  -- copies so poolChan is never read from
+  void $ forkIO $ forever $ readChan poolChan
+  metricsState <- Metrics.init (configDbPoolSize testCfg)
   pool <- P.acquire $ P.settings
     [ P.size 3
     , P.acquisitionTimeout 10
     , P.agingTimeout 60
     , P.idlenessTimeout 60
     , P.staticConnectionSettings (toUtf8 $ configDbUri testCfg)
+    -- make sure metrics are updated and pool observations published to poolChan
+    , P.observationHandler $ (writeChan poolChan <> Metrics.observationMetrics metricsState) . HasqlPoolObs
     ]
 
   actualPgVersion <- either (panic . show) id <$> P.use pool (queryPgVersion False)
@@ -86,7 +97,6 @@ main = do
   baseSchemaCache <- loadSCache pool testCfg
   sockets <- AppState.initSockets testCfg
   loggerState <- Logger.init
-  metricsState <- Metrics.init (configDbPoolSize testCfg)
 
   let
     initApp sCache st config = do
@@ -95,6 +105,15 @@ main = do
       AppState.putSchemaCache appState (Just sCache)
       return (st, postgrest (configLogLevel config) appState (pure ()))
 
+    initObservationsApp sCache config = do
+      -- duplicate poolChan as a starting point
+      obsChan <- dupChan poolChan
+      stateObsChan <- newObsChan obsChan
+      appState <- AppState.initWithPool sockets pool config loggerState metricsState (Metrics.observationMetrics metricsState <> writeChan obsChan)
+      AppState.putPgVersion appState actualPgVersion
+      AppState.putSchemaCache appState (Just sCache)
+      return (SpecState appState metricsState stateObsChan, postgrest (configLogLevel config) appState (pure ()))
+
     -- For tests that run with the same schema cache
     app = initApp baseSchemaCache ()
 
@@ -123,6 +142,7 @@ main = do
       obsApp               = app testObservabilityCfg
       serverTiming         = app testCfgServerTiming
       aggregatesEnabled    = app testCfgAggregatesEnabled
+      observationsApp      = initObservationsApp baseSchemaCache testCfg
 
       extraSearchPathApp   = appDbs testCfgExtraSearchPath
       unicodeApp           = appDbs testUnicodeCfg
@@ -278,6 +298,9 @@ main = do
     before (initApp baseSchemaCache metricsState testCfgJwtCache) $
       describe "Feature.Auth.JwtCacheSpec" Feature.Auth.JwtCacheSpec.spec
 
+    before observationsApp $
+      describe "Feature.MetricsSpec" Feature.MetricsSpec.spec
+
   where
     loadSCache pool conf =
       either (panic.show) id <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)
diff --git a/test/spec/SpecHelper.hs b/test/spec/SpecHelper.hs
@@ -1,10 +1,13 @@
 {-# LANGUAGE AllowAmbiguousTypes       #-}
+{-# LANGUAGE DeriveAnyClass            #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleContexts          #-}
+{-# LANGUAGE FlexibleInstances         #-}
 {-# LANGUAGE RankNTypes                #-}
 {-# LANGUAGE ScopedTypeVariables       #-}
 {-# LANGUAGE TupleSections             #-}
 {-# LANGUAGE TypeApplications          #-}
+{-# LANGUAGE TypeOperators             #-}
 module SpecHelper where
 
 import           Control.Lens           ((^?))
@@ -35,17 +38,25 @@ import Test.Hspec
 import Test.Hspec.Wai
 import Text.Heredoc
 
-import Data.String                       (String)
-import PostgREST.Config                  (AppConfig (..),
-                                          JSPathExp (..),
-                                          LogLevel (..),
-                                          OpenAPIMode (..),
-                                          Verbosity (..), parseSecret)
-import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
-import Prometheus                        (Counter, getCounter)
-import Protolude                         hiding (get, toS)
-import Protolude.Conv                    (toS)
-import Test.Hspec.Expectations.Contrib   (annotate)
+import qualified Data.List                         as DL
+import           Data.String                       (String)
+import qualified Data.Text                         as T
+import qualified PostgREST.AppState                as AppState
+import           PostgREST.Config                  (AppConfig (..),
+                                                    JSPathExp (..),
+                                                    LogLevel (..),
+                                                    OpenAPIMode (..),
+                                                    Verbosity (..),
+                                                    parseSecret)
+import qualified PostgREST.Metrics                 as Metrics
+import           PostgREST.Observation             (Observation (..))
+import           PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
+import           Prometheus                        (Counter,
+                                                    getCounter)
+import           Protolude                         hiding (get, toS)
+import           Protolude.Conv                    (toS)
+import           System.Timeout                    (timeout)
+import           Test.Hspec.Expectations.Contrib   (annotate)
 
 filterAndMatchCT :: BS.ByteString -> MatchHeader
 filterAndMatchCT val = MatchHeader $ \headers _ ->
@@ -381,3 +392,72 @@ expectCounter :: forall s st m. (KnownSymbol s, HasField s st Counter, MonadIO m
 expectCounter = expectField @s intCounter
   where
     intCounter = ((round @Double @Int) <$>) . getCounter
+
+data TimeoutException = TimeoutException deriving (Show, Exception)
+
+accumulateUntilTimeout :: Int -> (s -> a -> s) -> s -> IO a -> IO s
+accumulateUntilTimeout t f start act = do
+  tid <- myThreadId
+  -- mask to make sure TimeoutException is not thrown before starting the loop
+  mask $ \unmask -> do
+    -- start timeout thread unmasking exceptions
+    ttid <- forkIOWithUnmask ($ (threadDelay t *> throwTo tid TimeoutException))
+    -- unmask effect
+    unmask (fix (\loop accum -> (act >>= loop . f accum) `onTimeout` pure accum) start)
+      -- make sure we catch timeout if happens before entering the loop
+      `onTimeout` pure start
+      -- make sure timer thread is killed on other exceptions
+      -- so that it won't throw TimeoutException later
+      `onException` killThread ttid
+  where
+    onTimeout m a = m `catch` \TimeoutException -> a
+
+data ObsChan = ObsChan (Chan Observation) (Chan Observation)
+
+newObsChan :: Chan Observation -> IO ObsChan
+newObsChan = fmap <$> ObsChan <*> dupChan
+
+-- read messages from copy chan and once condition is met drain original to the same point
+-- upon timeout report error and messages remaining in the original chan
+-- that way we report messages since last successful read
+waitForObs :: HasCallStack => ObsChan -> Int -> Text -> (Observation -> Maybe a) -> IO ()
+waitForObs (ObsChan orig copy) t msg f =
+  timeout t (readUntil copy *> readUntil orig) >>= maybe failTimeout mempty
+  where
+    failTimeout = takeUntilTimeout decisecond (readChan orig)
+        >>= expectationFailure . DL.unlines . fmap show . (failureMessageHeader :) . fmap obsDiagMessage
+    failureMessageHeader = "Timeout waiting for " <> msg <> " at " <> loc <> ". Remaining observations:"
+    readUntil = void . untilM (pure . not . null . f) . readChan
+    loc = fromMaybe "(unknown)" . head $ (T.pack . prettySrcLoc . snd <$> getCallStack callStack)
+    -- execute effectful computation until result meets provided condition
+    untilM cond m = fix $ \loop -> m >>= \a -> ifM (cond a) (pure a) loop
+    -- duplicate the provided channel and construct wairFor function binding both channels
+    -- accumulate effecful computation results into a list for specified time
+    takeUntilTimeout t' = fmap reverse . accumulateUntilTimeout t' (flip (:)) []
+    obsDiagMessage (HasqlPoolObs o) = show o
+    obsDiagMessage o@(DBListenStart host port name channel) = constrName o <> show (host, port, name, channel)
+    obsDiagMessage o = constrName o
+    decisecond = 100000
+
+data SpecState = SpecState {
+  specAppState :: AppState.AppState,
+  specMetrics  :: Metrics.MetricsState,
+  specObsChan  :: ObsChan
+}
+
+-- helpers used to produce observation diagnostics in waitForObs
+constrName :: (HasConstructor (Rep a), Generic a)=> a -> Text
+constrName = genericConstrName . from
+
+class HasConstructor f where
+  genericConstrName :: f x -> Text
+
+instance HasConstructor f => HasConstructor (D1 c f) where
+  genericConstrName (M1 x) = genericConstrName x
+
+instance (HasConstructor x, HasConstructor y) => HasConstructor (x :+: y) where
+  genericConstrName (L1 l) = genericConstrName l
+  genericConstrName (R1 r) = genericConstrName r
+
+instance Constructor c => HasConstructor (C1 c f) where
+  genericConstrName = T.pack . conName
