packages feed

aws-xray-client-persistent (empty) → 0.1.0.0

raw patch · 4 files changed

+225/−0 lines, 4 filesdep +aws-xray-clientdep +basedep +conduit

Dependencies added: aws-xray-client, base, conduit, containers, lens, persistent, random, resourcet, text, time

Files

+ CHANGELOG.md view
@@ -0,0 +1,7 @@+## [*Unreleased*](https://github.com/freckle/aws-xray-client/tree/aws-xray-client-persistent-v0.1.0.0...main)++None++## [v0.1.0.0](https://github.com/freckle/aws-xray-client/tree/aws-xray-client-persistent-v0.1.0.0)++First tagged release.
+ README.md view
@@ -0,0 +1,54 @@+# aws-xray-client-persistent++A Haskell client for integrating AWS X-Ray with [Persistent](https://hackage.haskell.org/package/persistent).++To use this, you may want choose to define a helper to annotate traces++```hs++-- ...++import Data.Pool+import Database.Persist.Sql+import Network.AWS.XRayClient.Persistent++-- ...++runSqlPoolXRay+  :: (backend ~ SqlBackend, MonadUnliftIO m)+  => Text+  -- ^ Subsegment name+  --+  -- The top-level subsegment will be named @\"<this> runSqlPool\"@ and the,+  -- with a lower-level subsegment named @\"<this> query\"@.+  --+  -> XRayVaultData -- ^ Vault data to trace with+  -> ReaderT backend m a+  -> Pool backend+  -> m a+runSqlPoolXRay name vaultData action pool =+  traceXRaySubsegment' vaultData (name <> " runSqlPool") id+    $ withRunInIO+    $ \run -> withResource pool $ \backend -> do+        let+          sendTrace = atomicallyAddVaultDataSubsegment vaultData+          stdGenIORef = xrayVaultDataStdGen vaultData+          subsegmentName = name <> " query"+        run . runSqlConn action =<< liftIO+          (xraySqlBackend sendTrace stdGenIORef subsegmentName backend)+```++Then you can use this in your `runDB` definition to trace DB calls++```hs+-- ...+instance YesodPersist App where+  type YesodPersistBackend App = SqlBackend+  runDB action = do+    pool <- getsYesod appPool+    mVaultData <- vaultDataFromRequest <$> waiRequest+      +    -- ...++    maybe runSqlPool (runSqlPoolXRay "runDB") mVaultData action' pool+```
+ aws-xray-client-persistent.cabal view
@@ -0,0 +1,47 @@+cabal-version:      1.12+name:               aws-xray-client-persistent+version:            0.1.0.0+license:            MIT+copyright:          2021 Renaissance Learning Inc+maintainer:         engineering@freckle.com+author:             Freckle R&D+homepage:           https://github.com/freckle/aws-xray-client#readme+bug-reports:        https://github.com/freckle/aws-xray-client/issues+synopsis:           A client for AWS X-Ray integration with Persistent.+description:+    Works with `aws-xray-client` to enable X-Ray tracing with Persistent.++build-type:         Simple+extra-source-files:+    README.md+    CHANGELOG.md++source-repository head+    type:     git+    location: https://github.com/freckle/aws-xray-client++library+    exposed-modules:    Network.AWS.XRayClient.Persistent+    hs-source-dirs:     library+    other-modules:      Paths_aws_xray_client_persistent+    default-language:   Haskell2010+    default-extensions:+        BangPatterns DataKinds DeriveAnyClass DeriveFoldable DeriveFunctor+        DeriveGeneric DeriveLift DeriveTraversable DerivingStrategies+        FlexibleContexts FlexibleInstances GADTs GeneralizedNewtypeDeriving+        LambdaCase MultiParamTypeClasses NoImplicitPrelude+        NoMonomorphismRestriction OverloadedStrings RankNTypes+        RecordWildCards ScopedTypeVariables StandaloneDeriving+        TypeApplications TypeFamilies++    build-depends:+        aws-xray-client >=0.1.0.0,+        base >=4.12.0.0 && <5,+        conduit >=1.3.1.1,+        containers >=0.6.0.1,+        lens >=4.17.1,+        persistent >=2.9.2,+        random >=1.1,+        resourcet >=1.2.2,+        text >=1.2.3.1,+        time >=1.8.0.2
+ library/Network/AWS/XRayClient/Persistent.hs view
@@ -0,0 +1,117 @@+module Network.AWS.XRayClient.Persistent+  ( xraySqlBackend+  ) where++import Prelude++import Conduit+import Control.Lens+import Control.Monad (void)+import Data.Acquire (Acquire)+import Data.Foldable (for_)+import Data.IORef+import qualified Data.Map as Map+import Data.Text (Text)+import Data.Time.Clock.POSIX+import Database.Persist+import Database.Persist.Sql+import Database.Persist.Sql.Types.Internal (IsPersistBackend(mkPersistBackend))+import Network.AWS.XRayClient.Segment+import Network.AWS.XRayClient.TraceId+import System.Random+import System.Random.XRayCustom++-- | Modify a SqlBackend to send trace data to X-Ray.+--+-- >>> runSqlConn sql (xraySqlBackend "my-query" sendToDaemon backend)+xraySqlBackend+  :: (IsPersistBackend backend, BaseBackend backend ~ SqlBackend)+  => (XRaySegment -> IO ())+  -> IORef StdGen+  -> Text+  -> backend+  -> IO backend+xraySqlBackend sendTrace stdGenIORef subsegmentName =+  fmap mkPersistBackend . modifyBackend . persistBackend+ where+  modifyBackend backend = do+    -- N.B. by default persistent caches a Map Text Statement for each+    -- SqlBackend, where Text is a SQL query. When we wrap a backend to run it+    -- with XRay, we have to modify each Statement to record query timing. If+    -- backends are long-lived, then this poses a problem because we will+    -- continually wrap the same Statement. Therefore, we clear this cache each+    -- time we want to monitor things with XRay.+    newConnStmtMap <- newIORef Map.empty+    pure backend+      { connPrepare = connPrepare' (connPrepare backend)+      , connBegin = binaryTimerWrapper "BEGIN" (connBegin backend)+      , connCommit = unaryTimerWrapper "COMMIT" (connCommit backend)+      , connRollback = unaryTimerWrapper "ROLLBACK" (connRollback backend)+      , connStmtMap = newConnStmtMap+      }++  connPrepare' baseConnPrepare sql = do+    -- Create an IORef to store the start time. This is populated when a query+    -- begins in 'stmtQuery', and is then used in stmtReset to compute the+    -- total time.+    startTimeIORef <- newIORef Nothing++    statement <- baseConnPrepare sql+    pure statement+      { stmtQuery = stmtQuery' statement startTimeIORef+      , stmtReset = stmtReset' statement startTimeIORef sql+      }++  stmtQuery'+    :: forall m+     . MonadIO m+    => Statement+    -> IORef (Maybe POSIXTime)+    -> [PersistValue]+    -> Acquire (ConduitT () [PersistValue] m ())+  stmtQuery' statement startTimeIORef vals = do+    -- Record start time in IORef+    liftIO $ getPOSIXTime >>= writeIORef startTimeIORef . Just++    -- Create the Source and return it+    stmtQuery statement vals++  stmtReset' :: Statement -> IORef (Maybe POSIXTime) -> Text -> IO ()+  stmtReset' statement startTimeIORef sql = do+    stmtReset statement++    -- If start time exists (it should) then send the trace+    mStartTime <- readIORef startTimeIORef+    for_ mStartTime $ \startTime ->+      sendQueryTrace sendTrace subsegmentName startTime stdGenIORef sql++  unaryTimerWrapper sql action x = do+    startTime <- getPOSIXTime+    result <- action x+    sendQueryTrace sendTrace sql startTime stdGenIORef sql+    pure result++  binaryTimerWrapper sql action x y = do+    startTime <- getPOSIXTime+    result <- action x y+    sendQueryTrace sendTrace sql startTime stdGenIORef sql+    pure result++sendQueryTrace+  :: (XRaySegment -> IO ())+  -> Text+  -> POSIXTime+  -> IORef StdGen+  -> Text+  -> IO ()+sendQueryTrace sendTrace subsegmentName startTime stdGenIORef sql = do+  -- Record end time+  endTime <- getPOSIXTime++  -- Generate trace and send it off+  segmentId <- withRandomGenIORef stdGenIORef generateXRaySegmentId+  void+    $ sendTrace+    $ xraySubsegment subsegmentName segmentId startTime (Just endTime)+    & xraySegmentSql+    ?~ (xraySegmentSqlDef & xraySegmentSqlSanitizedQuery ?~ sql)