-
-
Notifications
You must be signed in to change notification settings - Fork 31
/
Persistent.hs
132 lines (121 loc) · 5.76 KB
/
Persistent.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module OpenTelemetry.Instrumentation.Persistent
( wrapSqlBackend
) where
import OpenTelemetry.Trace.Core
import OpenTelemetry.Context
import Data.Acquire.Internal
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Database.Persist.Sql
import Database.Persist.SqlBackend (setConnHooks, emptySqlBackendHooks, MkSqlBackendArgs (connRDBMS), getRDBMS, getConnVault, modifyConnVault)
import Database.Persist.SqlBackend.Internal
import Control.Monad.IO.Class
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.Vault.Strict as Vault
import OpenTelemetry.Attributes (Attributes)
import OpenTelemetry.Resource
import UnliftIO.Exception
import OpenTelemetry.Trace.Monad (MonadTracer(..))
import Control.Monad.Reader
import qualified Data.Text as T
import OpenTelemetry.Context.ThreadLocal (getContext, adjustContext)
instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlBackend m) where
getTracer = lift OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlReadBackend m) where
getTracer = lift OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlWriteBackend m) where
getTracer = lift OpenTelemetry.Trace.Monad.getTracer
originalConnectionKey :: Vault.Key SqlBackend
originalConnectionKey = unsafePerformIO Vault.newKey
{-# NOINLINE originalConnectionKey #-}
insertOriginalConnection :: SqlBackend -> SqlBackend -> SqlBackend
insertOriginalConnection conn original = modifyConnVault (Vault.insert originalConnectionKey original) conn
lookupOriginalConnection :: SqlBackend -> Maybe SqlBackend
lookupOriginalConnection = Vault.lookup originalConnectionKey . getConnVault
connectionLevelAttributesKey :: Vault.Key [(Text, Attribute)]
connectionLevelAttributesKey = unsafePerformIO Vault.newKey
{-# NOINLINE connectionLevelAttributesKey #-}
-- | Wrap a 'SqlBackend' with appropriate tracing context and attributes
-- so that queries are tracked appropriately in the tracing hierarchy.
wrapSqlBackend
:: MonadIO m
=> [(Text, Attribute)]
-- ^ Attributes that are specific to providers like MySQL, PostgreSQL, etc.
-> SqlBackend
-> m SqlBackend
wrapSqlBackend attrs conn_ = do
tp <- getGlobalTracerProvider
let conn = Data.Maybe.fromMaybe conn_ (lookupOriginalConnection conn_)
-- TODO add schema to tracerOptions?
let t = makeTracer tp "hs-opentelemetry-persistent" tracerOptions
let hooks = emptySqlBackendHooks
{ hookGetStatement = \conn sql stmt -> do
pure $ Statement
{ stmtQuery = \ps -> do
ctxt <- getContext
let spanCreator = do
s <- createSpan
t
ctxt
sql
(defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute sql) : attrs })
adjustContext (insertSpan s)
pure (lookupSpan ctxt, s)
spanCleanup (parent, s) = do
s `endSpan` Nothing
adjustContext $ \ctx ->
maybe ctx (`insertSpan` ctx) parent
(p, child) <- mkAcquire spanCreator spanCleanup
annotateBasics child conn
case stmtQuery stmt ps of
Acquire stmtQueryAcquireF -> Acquire $ \f ->
handleAny
(\(SomeException err) -> do
recordException child [] Nothing err
endSpan child Nothing
throwIO err
)
(stmtQueryAcquireF f)
, stmtExecute = \ps -> do
inSpan' t sql (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute sql) : attrs }) $ \s -> do
annotateBasics s conn
stmtExecute stmt ps
, stmtReset = stmtReset stmt
, stmtFinalize = stmtFinalize stmt
}
}
let conn' = conn
{ connHooks = hooks
, connBegin = \f mIso -> do
let statement = "begin transaction" <> case mIso of
Nothing -> mempty
Just ReadUncommitted -> " isolation level read uncommitted"
Just ReadCommitted -> " isolation level read committed"
Just RepeatableRead -> " isolation level repeatable read"
Just Serializable -> " isolation level serializable"
let attrs' = ("db.statement", toAttribute statement) : attrs
inSpan' t statement (defaultSpanArguments { kind = Client, attributes = attrs' }) $ \s -> do
annotateBasics s conn
connBegin conn f mIso
, connCommit = \f -> do
inSpan' t "commit" (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute ("commit" :: Text)): attrs }) $ \s -> do
annotateBasics s conn
connCommit conn f
, connRollback = \f -> do
inSpan' t "rollback" (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute ("rollback" :: Text)): attrs }) $ \s -> do
annotateBasics s conn
connRollback conn f
, connClose = do
inSpan' t "close connection" (defaultSpanArguments { kind = Client, attributes = attrs }) $ \s -> do
annotateBasics s conn
connClose conn
}
pure $ insertOriginalConnection conn' conn
annotateBasics :: MonadIO m => Span -> SqlBackend -> m ()
annotateBasics span conn = do
addAttributes span
[ ("db.system", toAttribute $ getRDBMS conn)
]