feat: add server-trace-header config for tracing

CHANGELOG.md

@@ -20,6 +20,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
    + Allows doing an anti join, e.g. `/projects?select=*,clients(*)&clients=is.null`
    + Allows using or across related tables conditions
  - #1100, Customizable OpenAPI title - @AnthonyFisi
+ - #2506, Add `server-trace-header` for tracing HTTP requests.  - @steve-chavez
+   + When the client sends the request header specified in the config it will be included in the response headers.
 
 ### Fixed
 

postgrest.cabal

@@ -187,6 +187,7 @@ test-suite spec
                       Feature.CorsSpec
                       Feature.ExtraSearchPathSpec
                       Feature.LegacyGucsSpec
+                      Feature.ObservabilitySpec
                       Feature.OpenApi.DisabledOpenApiSpec
                       Feature.OpenApi.IgnorePrivOpenApiSpec
                       Feature.OpenApi.OpenApiSpec

src/PostgREST/App.hs

@@ -48,7 +48,7 @@ import PostgREST.ApiRequest       (Action (..), ApiRequest (..),
                                    Mutation (..), Target (..))
 import PostgREST.AppState         (AppState)
 import PostgREST.Auth             (AuthResult (..))
-import PostgREST.Config           (AppConfig (..), LogLevel (..))
+import PostgREST.Config           (AppConfig (..))
 import PostgREST.Config.PgVersion (PgVersion (..))
 import PostgREST.Error            (Error)
 import PostgREST.Query            (DbHandler)
@@ -73,7 +73,7 @@ run installHandlers maybeRunWithSocket appState = do
 
   Workers.runAdmin conf appState $ serverSettings conf
 
-  let app = postgrest configLogLevel appState (Workers.connectionWorker appState)
+  let app = postgrest conf appState (Workers.connectionWorker appState)
 
   case configServerUnixSocket of
     Just socket ->
@@ -97,17 +97,18 @@ serverSettings AppConfig{..} =
     & setServerName ("postgrest/" <> prettyVersion)
 
 -- | PostgREST application
-postgrest :: LogLevel -> AppState.AppState -> IO () -> Wai.Application
+postgrest :: AppConfig -> AppState.AppState -> IO () -> Wai.Application
-postgrest logLevel appState connWorker =
+postgrest conf appState connWorker =
+  Response.traceHeaderMiddleware conf .
   Cors.middleware .
   Auth.middleware appState .
-  Logger.middleware logLevel $
+  Logger.middleware (configLogLevel conf) $
     -- fromJust can be used, because the auth middleware will **always** add
     -- some AuthResult to the vault.
     \req respond -> case fromJust $ Auth.getResult req of
       Left err -> respond $ Error.errorResponseFor err
       Right authResult -> do
-        conf <- AppState.getConfig appState
+        appConf <- AppState.getConfig appState -- the config must be read again because it can reload
         maybeSchemaCache <- AppState.getSchemaCache appState
         pgVer <- AppState.getPgVersion appState
         jsonDbS <- AppState.getJsonDbS appState
@@ -115,7 +116,7 @@ postgrest logLevel appState connWorker =
         let
           eitherResponse :: IO (Either Error Wai.Response)
           eitherResponse =
-            runExceptT $ postgrestResponse appState conf maybeSchemaCache jsonDbS pgVer authResult req
+            runExceptT $ postgrestResponse appState appConf maybeSchemaCache jsonDbS pgVer authResult req
 
         response <- either Error.errorResponseFor identity <$> eitherResponse
         -- Launch the connWorker when the connection is down.  The postgrest

src/PostgREST/Config.hs

@@ -32,6 +32,7 @@ import qualified Data.Aeson             as JSON
 import qualified Data.ByteString        as BS
 import qualified Data.ByteString.Base64 as B64
 import qualified Data.ByteString.Lazy   as LBS
+import qualified Data.CaseInsensitive   as CI
 import qualified Data.Configurator      as C
 import qualified Data.Map.Strict        as M
 import qualified Data.Text              as T
@@ -93,6 +94,7 @@ data AppConfig = AppConfig
   , configRawMediaTypes            :: [MediaType]
   , configServerHost               :: Text
   , configServerPort               :: Int
+  , configServerTraceHeader        :: Maybe (CI.CI BS.ByteString)
   , configServerUnixSocket         :: Maybe FilePath
   , configServerUnixSocketMode     :: FileMode
   , configAdminServerPort          :: Maybe Int
@@ -150,6 +152,7 @@ toText conf =
       ,("raw-media-types",           q . T.decodeUtf8 . BS.intercalate "," . fmap toMime . configRawMediaTypes)
       ,("server-host",               q . configServerHost)
       ,("server-port",                   show . configServerPort)
+      ,("server-trace-header",       q . T.decodeUtf8 . maybe mempty CI.original . configServerTraceHeader)
       ,("server-unix-socket",        q . maybe mempty T.pack . configServerUnixSocket)
       ,("server-unix-socket-mode",   q . T.pack . showSocketMode)
       ,("admin-server-port",             maybe "\"\"" show . configAdminServerPort)
@@ -247,6 +250,7 @@ parser optPath env dbSettings =
     <*> (maybe [] (fmap (MTOther . encodeUtf8) . splitOnCommas) <$> optValue "raw-media-types")
     <*> (fromMaybe "!4" <$> optString "server-host")
     <*> (fromMaybe 3000 <$> optInt "server-port")
+    <*> (fmap (CI.mk . encodeUtf8) <$> optString "server-trace-header")
     <*> (fmap T.unpack <$> optString "server-unix-socket")
     <*> parseSocketFileMode "server-unix-socket-mode"
     <*> optInt "admin-server-port"

src/PostgREST/Response.hs

@@ -14,12 +14,14 @@ module PostgREST.Response
   , addRetryHint
   , isServiceUnavailable
   , optionalRollback
+  , traceHeaderMiddleware
   ) where
 
 import qualified Data.Aeson                as JSON
 import qualified Data.ByteString.Char8     as BS
 import qualified Data.ByteString.Lazy      as LBS
 import qualified Data.HashMap.Strict       as HM
+import qualified Data.List                 as L
 import           Data.Text.Read            (decimal)
 import qualified Network.HTTP.Types.Header as HTTP
 import qualified Network.HTTP.Types.Status as HTTP
@@ -291,3 +293,11 @@ addHeadersIfNotIncluded :: [HTTP.Header] -> [HTTP.Header] -> [HTTP.Header]
 addHeadersIfNotIncluded newHeaders initialHeaders =
   filter (\(nk, _) -> isNothing $ find (\(ik, _) -> ik == nk) initialHeaders) newHeaders ++
   initialHeaders
+
+traceHeaderMiddleware :: AppConfig -> Wai.Middleware
+traceHeaderMiddleware AppConfig{configServerTraceHeader} app req respond =
+  case configServerTraceHeader of
+    Nothing -> app req respond
+    Just hdr ->
+      let hdrVal = L.lookup hdr $ Wai.requestHeaders req in
+      app req (respond . Wai.mapResponseHeaders ([(hdr, fromMaybe mempty hdrVal)] ++))

test/io/configs/expected/aliases.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = ""
 raw-media-types = ""
 server-host = "!4"
 server-port = 3000
+server-trace-header = ""
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-port = ""

test/io/configs/expected/boolean-numeric.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = ""
 raw-media-types = ""
 server-host = "!4"
 server-port = 3000
+server-trace-header = ""
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-port = ""

test/io/configs/expected/boolean-string.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = ""
 raw-media-types = ""
 server-host = "!4"
 server-port = 3000
+server-trace-header = ""
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-port = ""

test/io/configs/expected/defaults.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = ""
 raw-media-types = ""
 server-host = "!4"
 server-port = 3000
+server-trace-header = ""
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-port = ""

test/io/configs/expected/no-defaults-with-db-other-authenticator.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = "https://otherexample.org/api"
 raw-media-types = "application/vnd.pgrst.other-db-config"
 server-host = "0.0.0.0"
 server-port = 80
+server-trace-header = "traceparent"
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-port = 3001

test/io/configs/expected/no-defaults-with-db.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = "https://example.org/api"
 raw-media-types = "application/vnd.pgrst.db-config"
 server-host = "0.0.0.0"
 server-port = 80
+server-trace-header = "CF-Ray"
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-port = 3001

test/io/configs/expected/no-defaults.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = "https://postgrest.org"
 raw-media-types = "application/vnd.pgrst.config"
 server-host = "0.0.0.0"
 server-port = 80
+server-trace-header = "X-Request-Id"
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-port = 3001

test/io/configs/expected/types.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = ""
 raw-media-types = ""
 server-host = "!4"
 server-port = 3000
+server-trace-header = ""
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-port = ""

test/io/configs/no-defaults-env.yaml

@@ -28,6 +28,7 @@ PGRST_OPENAPI_SERVER_PROXY_URI: 'https://postgrest.org'
 PGRST_RAW_MEDIA_TYPES: application/vnd.pgrst.config
 PGRST_SERVER_HOST: 0.0.0.0
 PGRST_SERVER_PORT: 80
+PGRST_SERVER_TRACE_HEADER: X-Request-Id
 PGRST_SERVER_UNIX_SOCKET: /tmp/pgrst_io_test.sock
 PGRST_SERVER_UNIX_SOCKET_MODE: 777
 PGRST_ADMIN_SERVER_PORT: 3001

test/io/configs/no-defaults.config

@@ -25,6 +25,7 @@ openapi-server-proxy-uri = "https://postgrest.org"
 raw-media-types = "application/vnd.pgrst.config"
 server-host = "0.0.0.0"
 server-port = 80
+server-trace-header = "X-Request-Id"
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-port = 3001

test/io/db_config.sql

@@ -17,6 +17,7 @@ ALTER ROLE db_config_authenticator SET pgrst.db_pre_request = 'test.custom_heade
 ALTER ROLE db_config_authenticator SET pgrst.db_max_rows = '1000';
 ALTER ROLE db_config_authenticator SET pgrst.db_extra_search_path = 'public, extensions';
 ALTER ROLE db_config_authenticator SET pgrst.not_existing = 'should be ignored';
+ALTER ROLE db_config_authenticator SET pgrst.server_trace_header = 'CF-Ray';
 
 -- override with database specific setting
 ALTER ROLE db_config_authenticator IN DATABASE :DBNAME SET pgrst.jwt_secret = 'OVERRIDE=REALLY=REALLY=REALLY=REALLY=VERY=SAFE';
@@ -60,6 +61,7 @@ ALTER ROLE other_authenticator SET pgrst.db_max_rows = '100';
 ALTER ROLE other_authenticator SET pgrst.db_extra_search_path = 'public, extensions, other';
 ALTER ROLE other_authenticator SET pgrst.openapi_mode = 'disabled';
 ALTER ROLE other_authenticator SET pgrst.openapi_security_active = 'false';
+ALTER ROLE other_authenticator SET pgrst.server_trace_header = 'traceparent';
 
 -- authenticator used for tests that manipulate statement timeout
 CREATE ROLE timeout_authenticator LOGIN NOINHERIT;

test/spec/Feature/ObservabilitySpec.hs

@@ -0,0 +1,34 @@
+module Feature.ObservabilitySpec where
+
+import Network.Wai (Application)
+
+import Network.HTTP.Types
+import Test.Hspec
+import Test.Hspec.Wai
+
+import Protolude
+
+spec :: SpecWith ((), Application)
+spec =
+  describe "Observability" $ do
+    it "includes the server trace header on the response" $ do
+      request methodHead "/"
+          [ ("X-Request-Id", "1") ]
+          ""
+        `shouldRespondWith`
+          ""
+          { matchHeaders = [ "X-Request-Id" <:> "1"] }
+
+      request methodHead "/projects"
+          [ ("X-Request-Id", "2") ]
+          ""
+        `shouldRespondWith`
+          ""
+          { matchHeaders = [ "X-Request-Id" <:> "2"] }
+
+      request methodHead "/rpc/add_them?a=2&b=4"
+          [ ("X-Request-Id", "3") ]
+          ""
+        `shouldRespondWith`
+          ""
+          { matchHeaders = [ "X-Request-Id" <:> "3"] }

test/spec/Main.hs

@@ -10,7 +10,7 @@ import Data.List.NonEmpty (toList)
 import Test.Hspec
 
 import PostgREST.App             (postgrest)
-import PostgREST.Config          (AppConfig (..), LogLevel (..))
+import PostgREST.Config          (AppConfig (..))
 import PostgREST.Config.Database (queryPgVersion)
 import PostgREST.SchemaCache     (querySchemaCache)
 import Protolude                 hiding (toList, toS)
@@ -29,6 +29,7 @@ import qualified Feature.ConcurrentSpec
 import qualified Feature.CorsSpec
 import qualified Feature.ExtraSearchPathSpec
 import qualified Feature.LegacyGucsSpec
+import qualified Feature.ObservabilitySpec
 import qualified Feature.OpenApi.DisabledOpenApiSpec
 import qualified Feature.OpenApi.IgnorePrivOpenApiSpec
 import qualified Feature.OpenApi.OpenApiSpec
@@ -83,7 +84,7 @@ main = do
       AppState.putSchemaCache appState (Just baseSchemaCache)
       when (isJust $ configDbRootSpec config) $
         AppState.putJsonDbS appState $ toS $ JSON.encode baseSchemaCache
-      return ((), postgrest LogCrit appState $ pure ())
+      return ((), postgrest config appState $ pure ())
 
     -- For tests that run with a different SchemaCache(depends on configSchemas)
     appDbs config = do
@@ -96,7 +97,7 @@ main = do
       AppState.putSchemaCache appState (Just customSchemaCache)
       when (isJust $ configDbRootSpec config) $
         AppState.putJsonDbS appState $ toS $ JSON.encode baseSchemaCache
-      return ((), postgrest LogCrit appState $ pure ())
+      return ((), postgrest config appState $ pure ())
 
   let withApp              = app testCfg
       maxRowsApp           = app testMaxRowsCfg
@@ -117,6 +118,7 @@ main = do
       testCfgLegacyGucsApp = app testCfgLegacyGucs
       planEnabledApp       = app testPlanEnabledCfg
       pgSafeUpdateApp      = app testPgSafeUpdateEnabledCfg
+      obsApp               = app testObservabilityCfg
 
       extraSearchPathApp   = appDbs testCfgExtraSearchPath
       unicodeApp           = appDbs testUnicodeCfg
@@ -247,6 +249,10 @@ main = do
     parallel $ before pgSafeUpdateApp $
       describe "Feature.Query.PgSafeUpdateSpec.spec" Feature.Query.PgSafeUpdateSpec.spec
 
+    -- this test runs with server-trace-header set
+    parallel $ before obsApp $
+      describe "Feature.ObservabilitySpec.spec" Feature.ObservabilitySpec.spec
+
     -- Note: the rollback tests can not run in parallel, because they test persistance and
     -- this results in race conditions
 

test/spec/SpecHelper.hs

@@ -10,7 +10,7 @@ import           Data.Scientific        (toRealFloat)
 import qualified Data.Set               as S
 
 import Data.Aeson           (Value (..), decode, encode)
-import Data.CaseInsensitive (CI (..), original)
+import Data.CaseInsensitive (CI (..), mk, original)
 import Data.List            (lookup)
 import Data.List.NonEmpty   (fromList)
 import Network.Wai.Test     (SResponse (simpleBody, simpleHeaders, simpleStatus))
@@ -103,6 +103,7 @@ baseCfg = let secret = Just $ encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configRawMediaTypes         = []
   , configServerHost            = "localhost"
   , configServerPort            = 3000
+  , configServerTraceHeader     = Nothing
   , configServerUnixSocket      = Nothing
   , configServerUnixSocketMode  = 432
   , configDbTxAllowOverride     = True
@@ -203,6 +204,9 @@ testCfgLegacyGucs = baseCfg { configDbUseLegacyGucs = False }
 testPgSafeUpdateEnabledCfg :: AppConfig
 testPgSafeUpdateEnabledCfg = baseCfg { configDbPreRequest = Just $ QualifiedIdentifier "test" "load_safeupdate" }
 
+testObservabilityCfg :: AppConfig
+testObservabilityCfg = baseCfg { configServerTraceHeader = Just $ mk "X-Request-Id" }
+
 analyzeTable :: Text -> IO ()
 analyzeTable tableName =
   void $ readProcess "psql" ["--set", "ON_ERROR_STOP=1", "-a", "-c", toS $ "ANALYZE test.\"" <> tableName <> "\""] []