scratch

PostgREST · Jan 29, 2024 · 64a0ee9 · 64a0ee9
1 parent bbc0bda
commit 64a0ee9
Show file tree

Hide file tree

Showing 5 changed files with 190 additions and 49 deletions.
diff --git a/postgrest.cabal b/postgrest.cabal
@@ -102,6 +102,7 @@ library
                     , hasql-transaction         >= 1.0.1 && < 1.1
                     , heredoc                   >= 0.2 && < 0.3
                     , http-types                >= 0.12.2 && < 0.13
+                    , hs-opentelemetry-sdk      >= 0.0.3.6 && < 0.0.4
                     , insert-ordered-containers >= 0.2.2 && < 0.3
                     , interpolatedstring-perl6  >= 1 && < 1.1
                     , jose                      >= 0.8.5.1 && < 0.12

diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
@@ -49,7 +49,7 @@ import PostgREST.AppState             (AppState)
 import PostgREST.Auth                 (AuthResult (..))
 import PostgREST.Config               (AppConfig (..))
 import PostgREST.Config.PgVersion     (PgVersion (..))
-import PostgREST.Error                (Error)
+import PostgREST.Error                (Error (..))
 import PostgREST.Query                (DbHandler)
 import PostgREST.Response.Performance (ServerTiming (..),
                                        serverTimingHeader)
@@ -63,6 +63,7 @@ import qualified Network.HTTP.Types    as HTTP
 import qualified Network.Socket        as NS
 import           Protolude             hiding (Handler)
 import           System.TimeIt         (timeItT)
+import OpenTelemetry.Trace (getGlobalTracerProvider, defaultSpanArguments, makeTracer, InstrumentationLibrary (..), TracerOptions (..), Tracer, SpanKind (..), SpanArguments(..), inSpan', addEvent, NewEvent (NewEvent, newEventAttributes, newEventTimestamp, newEventName), setStatus, SpanStatus(..), Span, ToAttribute (toAttribute))
 
 type Handler = ExceptT Error
 
@@ -136,6 +137,7 @@ postgrestResponse
   -> Wai.Request
   -> Handler IO Wai.Response
 postgrestResponse appState conf@AppConfig{..} maybeSchemaCache pgVer authResult@AuthResult{..} req = do
+  t <- liftIO getTracer
   sCache <-
     case maybeSchemaCache of
       Just sCache ->
@@ -146,12 +148,20 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache pgVer authResult@
   body <- lift $ Wai.strictRequestBody req
 
   (parseTime, apiRequest) <-
-    calcTiming configServerTimingEnabled $
+    calcTiming' configServerTimingEnabled $
       liftEither . mapLeft Error.ApiRequestError $
         ApiRequest.userApiRequest conf req body sCache
 
   let jwtTime = if configServerTimingEnabled then Auth.getJwtDur req else Nothing
-  handleRequest authResult conf appState (Just authRole /= configDbAnonRole) configDbPreparedStatements pgVer apiRequest sCache jwtTime parseTime
+  liftIO $ inSpan' t "handleRequest" defaultSpanArguments {kind = Server} $ \span -> do
+    result <- runExceptT $ handleRequest span authResult conf appState (Just authRole /= configDbAnonRole) configDbPreparedStatements pgVer apiRequest sCache jwtTime parseTime
+    case result of
+      Left err -> do
+        setStatus span (Error "Boo!")
+        inSpan' t "errorResponseFor" defaultSpanArguments {kind = Server} $ \span' -> do
+          _ <- addEvent span' $ NewEvent { newEventName = "error", newEventAttributes = HM.fromList [], newEventTimestamp = Nothing }
+          return $ Error.errorResponseFor err
+      Right resp -> return resp
 
 runDbHandler :: AppState.AppState -> AppConfig -> SQL.IsolationLevel -> SQL.Mode -> Bool -> Bool -> DbHandler b -> Handler IO b
 runDbHandler appState config isoLvl mode authenticated prepared handler = do
@@ -165,62 +175,62 @@ runDbHandler appState config isoLvl mode authenticated prepared handler = do
 
   liftEither resp
 
-handleRequest :: AuthResult -> AppConfig -> AppState.AppState -> Bool -> Bool -> PgVersion -> ApiRequest -> SchemaCache -> Maybe Double -> Maybe Double -> Handler IO Wai.Response
-handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@ApiRequest{..} sCache jwtTime parseTime =
+handleRequest :: Span -> AuthResult -> AppConfig -> AppState.AppState -> Bool -> Bool -> PgVersion -> ApiRequest -> SchemaCache -> Maybe Double -> Maybe Double -> Handler IO Wai.Response
+handleRequest span AuthResult{..} conf appState authenticated prepared pgVer apiReq@ApiRequest{..} sCache jwtTime parseTime =
   case (iAction, iTarget) of
     (ActionRead headersOnly, TargetIdent identifier) -> do
-      (planTime', wrPlan) <- withTiming $ liftEither $ Plan.wrappedReadPlan identifier conf sCache apiReq
-      (txTime', resultSet) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.wrTxMode wrPlan) $ Query.readQuery wrPlan conf apiReq
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.readResponse wrPlan headersOnly identifier apiReq resultSet
+      (planTime', wrPlan) <- withTiming "queryTime" $ liftEither $ Plan.wrappedReadPlan identifier conf sCache apiReq
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.wrTxMode wrPlan) $ Query.readQuery wrPlan conf apiReq
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.readResponse wrPlan headersOnly identifier apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionMutate MutationCreate, TargetIdent identifier) -> do
-      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationCreate apiReq identifier conf sCache
-      (txTime', resultSet) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.createQuery mrPlan apiReq conf
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.createResponse identifier mrPlan apiReq resultSet
+      (planTime', mrPlan) <- withTiming "queryTime" $ liftEither $ Plan.mutateReadPlan MutationCreate apiReq identifier conf sCache
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.createQuery mrPlan apiReq conf
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.createResponse identifier mrPlan apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionMutate MutationUpdate, TargetIdent identifier) -> do
-      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationUpdate apiReq identifier conf sCache
-      (txTime', resultSet) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.updateQuery mrPlan apiReq conf
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.updateResponse mrPlan apiReq resultSet
+      (planTime', mrPlan) <- withTiming "queryTime" $ liftEither $ Plan.mutateReadPlan MutationUpdate apiReq identifier conf sCache
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.updateQuery mrPlan apiReq conf
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.updateResponse mrPlan apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionMutate MutationSingleUpsert, TargetIdent identifier) -> do
-      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationSingleUpsert apiReq identifier conf sCache
-      (txTime', resultSet) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.singleUpsertQuery mrPlan apiReq conf
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.singleUpsertResponse mrPlan apiReq resultSet
+      (planTime', mrPlan) <- withTiming "queryTime" $ liftEither $ Plan.mutateReadPlan MutationSingleUpsert apiReq identifier conf sCache
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.singleUpsertQuery mrPlan apiReq conf
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.singleUpsertResponse mrPlan apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionMutate MutationDelete, TargetIdent identifier) -> do
-      (planTime', mrPlan) <- withTiming $ liftEither $ Plan.mutateReadPlan MutationDelete apiReq identifier conf sCache
-      (txTime', resultSet) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.deleteQuery mrPlan apiReq conf
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.deleteResponse mrPlan apiReq resultSet
+      (planTime', mrPlan) <- withTiming "queryTime" $ liftEither $ Plan.mutateReadPlan MutationDelete apiReq identifier conf sCache
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.mrTxMode mrPlan) $ Query.deleteQuery mrPlan apiReq conf
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.deleteResponse mrPlan apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionInvoke invMethod, TargetProc identifier _) -> do
-      (planTime', cPlan) <- withTiming $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq invMethod
-      (txTime', resultSet) <- withTiming $ runQuery (fromMaybe roleIsoLvl $ pdIsoLvl (Plan.crProc cPlan)) (pdTimeout $ Plan.crProc cPlan) (Plan.crTxMode cPlan) $ Query.invokeQuery (Plan.crProc cPlan) cPlan apiReq conf pgVer
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.invokeResponse cPlan invMethod (Plan.crProc cPlan) apiReq resultSet
+      (planTime', cPlan) <- withTiming "queryTime" $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq invMethod
+      (txTime', resultSet) <- withTiming "planTime" $ runQuery (fromMaybe roleIsoLvl $ pdIsoLvl (Plan.crProc cPlan)) (pdTimeout $ Plan.crProc cPlan) (Plan.crTxMode cPlan) $ Query.invokeQuery (Plan.crProc cPlan) cPlan apiReq conf pgVer
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.invokeResponse cPlan invMethod (Plan.crProc cPlan) apiReq resultSet
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionInspect headersOnly, TargetDefaultSpec tSchema) -> do
-      (planTime', iPlan) <- withTiming $ liftEither $ Plan.inspectPlan apiReq
-      (txTime', oaiResult) <- withTiming $ runQuery roleIsoLvl Nothing (Plan.ipTxmode iPlan) $ Query.openApiQuery sCache pgVer conf tSchema
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.openApiResponse (T.decodeUtf8 prettyVersion, docsVersion) headersOnly oaiResult conf sCache iSchema iNegotiatedByProfile
+      (planTime', iPlan) <- withTiming "queryTime" $ liftEither $ Plan.inspectPlan apiReq
+      (txTime', oaiResult) <- withTiming "planTime" $ runQuery roleIsoLvl Nothing (Plan.ipTxmode iPlan) $ Query.openApiQuery sCache pgVer conf tSchema
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.openApiResponse (T.decodeUtf8 prettyVersion, docsVersion) headersOnly oaiResult conf sCache iSchema iNegotiatedByProfile
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' txTime' respTime') pgrst
 
     (ActionInfo, TargetIdent identifier) -> do
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.infoIdentResponse identifier sCache
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.infoIdentResponse identifier sCache
       return $ pgrstResponse (ServerTiming jwtTime parseTime Nothing Nothing respTime') pgrst
 
     (ActionInfo, TargetProc identifier _) -> do
-      (planTime', cPlan) <- withTiming $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq ApiRequest.InvHead
-      (respTime', pgrst) <- withTiming $ liftEither $ Response.infoProcResponse (Plan.crProc cPlan)
+      (planTime', cPlan) <- withTiming "planTime" $ liftEither $ Plan.callReadPlan identifier conf sCache apiReq ApiRequest.InvHead
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither $ Response.infoProcResponse (Plan.crProc cPlan)
       return $ pgrstResponse (ServerTiming jwtTime parseTime planTime' Nothing respTime') pgrst
 
     (ActionInfo, TargetDefaultSpec _) -> do
-      (respTime', pgrst) <- withTiming $ liftEither Response.infoRootResponse
+      (respTime', pgrst) <- withTiming "responseTime" $ liftEither Response.infoRootResponse
       return $ pgrstResponse (ServerTiming jwtTime parseTime Nothing Nothing respTime') pgrst
 
     _ ->
@@ -239,10 +249,20 @@ handleRequest AuthResult{..} conf appState authenticated prepared pgVer apiReq@A
     pgrstResponse :: ServerTiming -> Response.PgrstResponse -> Wai.Response
     pgrstResponse timing (Response.PgrstResponse st hdrs bod) = Wai.responseLBS st (hdrs ++ ([serverTimingHeader timing | configServerTimingEnabled conf])) bod
 
-    withTiming = calcTiming $ configServerTimingEnabled conf
+    withTiming label = calcTiming span label $ configServerTimingEnabled conf
 
-calcTiming :: Bool -> Handler IO a -> Handler IO (Maybe Double, a)
-calcTiming timingEnabled f = if timingEnabled
+calcTiming :: Span -> Text -> Bool -> Handler IO a -> Handler IO (Maybe Double, a)
+calcTiming span label timingEnabled f = do 
+  (t, r) <- timeItT f
+  _ <- addEvent span $ NewEvent { newEventName = label, newEventAttributes = HM.fromList [("time", toAttribute t)], newEventTimestamp = Nothing }
+  if timingEnabled
+    then do
+      pure (Just t, r)
+    else do
+      pure (Nothing, r)
+
+calcTiming' :: Bool -> Handler IO a -> Handler IO (Maybe Double, a)
+calcTiming' timingEnabled f = if timingEnabled
     then do
       (t, r) <- timeItT f
       pure (Just t, r)
@@ -265,3 +285,11 @@ addRetryHint delay response = do
 
 isServiceUnavailable :: Wai.Response -> Bool
 isServiceUnavailable response = Wai.responseStatus response == HTTP.status503
+
+getTracer :: IO Tracer 
+getTracer = do
+  tp <- getGlobalTracerProvider
+  return $ 
+    makeTracer tp 
+      InstrumentationLibrary { libraryVersion = decodeUtf8 prettyVersion, libraryName = "postgrest"} 
+      TracerOptions {tracerSchema=Nothing}
diff --git a/src/PostgREST/CLI.hs b/src/PostgREST/CLI.hs
@@ -28,9 +28,10 @@ import qualified PostgREST.Config   as Config
 
 import Protolude hiding (hPutStrLn)
 
+import OpenTelemetry.Trace
 
 main :: CLI -> IO ()
-main CLI{cliCommand, cliPath} = do
+main CLI{cliCommand, cliPath} = withTracer $ \_tracer -> do
   conf@AppConfig{..} <-
     either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
 
@@ -46,6 +47,12 @@ main CLI{cliCommand, cliPath} = do
         putStr . Config.toText =<< AppState.getConfig appState
       CmdDumpSchema -> putStrLn =<< dumpSchema appState
       CmdRun -> App.run appState)
+  where
+    withTracer :: ((TracerOptions -> Tracer) -> IO c) -> IO c 
+    withTracer f = bracket 
+      initializeGlobalTracerProvider
+      shutdownTracerProvider
+      (\tracerProvider -> f $ makeTracer tracerProvider "PostgREST")
 
 -- | Dump SchemaCache schema to JSON
 dumpSchema :: AppState -> IO LBS.ByteString

diff --git a/stack.yaml b/stack.yaml
@@ -16,3 +16,20 @@ extra-deps:
   - hasql-pool-0.10
   - megaparsec-9.2.2
   - postgresql-libpq-0.10.0.0
+  # - hs-opentelemetry-sdk-0.0.3.6
+  # - hs-opentelemetry-api-0.1.0.0@sha256:8af01d0c81dd1af6d3293b105178fd9bfa0057c9eb88ac24d3c440bff660abe3,3705
+  # - hs-opentelemetry-exporter-otlp-0.0.1.5@sha256:89b0a6481096a338fa6383fbdf08ccaa0eb7bb009c4cbb340894eac33e55c5de,2214
+  # - hs-opentelemetry-propagator-b3-0.0.1.1@sha256:f0e9da77a888b89f81e5f5186788d5ace2a665e8914f6b446712a1c2edf17743,1854
+  # - hs-opentelemetry-propagator-w3c-0.0.1.3@sha256:5dc2dbdd6b0a4e434ca5fd949e9ebe5611a5d513ef58009b935e9e810cc85d1b,1852
+  # - hs-opentelemetry-otlp-0.0.1.0@sha256:88bb6b68f172a336f78018b0823f47363fb7408eb19f7301489f81ad4d5c0f33,2307
+  - github: 'iand675/hs-opentelemetry'
+    commit: '54251bd15565e3a07cda0cecfa544d209ae63df6'
+    subdirs:
+      - sdk
+      - api
+      - 'exporters/otlp'
+      - 'propagators/b3'
+      - 'propagators/w3c'
+      - otlp
+  - thread-utils-context-0.3.0.4@sha256:e763da1c6cab3b6d378fb670ca74aa9bf03c9b61b6fcf7628c56363fb0e3e71e,1671
+  - thread-utils-finalizers-0.1.1.0@sha256:24944b71d9f1d01695a5908b4a3b44838fab870883114a323336d537995e0a5b,1381