content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE OverloadedLists #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeApplications #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module OpenTelemetry.Instrumentation.Persistent (
    wrapSqlBackend,
    wrapSqlBackend',
  ) where
  
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.Monad.Reader
  import Data.Acquire.Internal
  import qualified Data.HashMap.Strict as H
  import Data.IORef
  import Data.Maybe (fromMaybe)
  import Data.Text (Text)
  import qualified Data.Text as T
  import qualified Data.Vault.Strict as Vault
  import Database.Persist.Sql (IsolationLevel (..), SqlReadBackend, SqlWriteBackend, Statement (..))
  import Database.Persist.SqlBackend (MkSqlBackendArgs (connRDBMS), emptySqlBackendHooks, getConnVault, getRDBMS, modifyConnVault, setConnHooks)
  import Database.Persist.SqlBackend.Internal
  import OpenTelemetry.Attributes (Attributes)
  import OpenTelemetry.Attributes.Map (AttributeMap)
  import OpenTelemetry.Common
  import OpenTelemetry.Context
  import OpenTelemetry.Context.ThreadLocal (adjustContext, getContext)
  import OpenTelemetry.Resource
  import OpenTelemetry.Trace.Core
  import OpenTelemetry.Trace.Monad (MonadTracer (..))
  import System.Clock
  import System.IO.Unsafe (unsafePerformIO)
  import UnliftIO.Exception
  
  
  {-
  Design notes:
  
  In some OTel export destinations like Honeycomb, the cost is per-span. Consquently, we want to minimize the number of spans we create. In particular, we want to avoid creating a span for every query, since they add up cost-wise.
  
  However, we also want to capture transactions as spans. Therefore, for pool acquisitions we track the time between trying to acquire the connection and the time the connection is obtained as an attribute on the initial span.
  -}
  
  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 AttributeMap
  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
    => AttributeMap
    -- ^ Attributes that are specific to providers like MySQL, PostgreSQL, etc.
    -> SqlBackend
    -> m SqlBackend
  wrapSqlBackend attrs conn_ = do
    tp <- getGlobalTracerProvider
    wrapSqlBackend' tp attrs conn_
  
  
  {- | Wrap a 'SqlBackend' with appropriate tracing context and attributes
  so that queries are tracked appropriately in the tracing hierarchy.
  -}
  wrapSqlBackend'
    :: MonadIO m
    => TracerProvider
    -> AttributeMap
    -- ^ Attributes that are specific to providers like MySQL, PostgreSQL, etc.
    -> SqlBackend
    -> m SqlBackend
  wrapSqlBackend' tp attrs conn_ = do
    let conn = Data.Maybe.fromMaybe conn_ (lookupOriginalConnection conn_)
    {- A connection is acquired when the connection pool is asked for a connection. The runSqlPool function in Persistent then
      immediately begins a transaction and ensures the transaction is committed or rolled back. Since we want to capture the
      transaction as a span, we have to use track the current Span in flight. We do this because we can't hand off
      the Span between connBegin/connCommit/connRollback as return values.
    -}
    connParentSpan <- liftIO $ newIORef Nothing
    connSpanInFlight <- liftIO $ newIORef Nothing
    -- TODO add schema to tracerOptions?
    let t = makeTracer tp $detectInstrumentationLibrary 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 = H.insert "db.statement" (toAttribute sql) attrs})
                              adjustContext (insertSpan s)
                              pure (lookupSpan ctxt, s)
                            spanCleanup (parent, s) = do
                              s `endSpan` Nothing
                              adjustContext $ \ctx ->
                                maybe (removeSpan 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 [("exception.escaped", toAttribute True)] Nothing err
                                  endSpan child Nothing
                                  throwIO err
                              )
                              (stmtQueryAcquireF f)
                    , stmtExecute = \ps -> do
                        inSpan' t sql (defaultSpanArguments {kind = Client, attributes = H.insert "db.statement" (toAttribute sql) attrs}) $ \s -> do
                          annotateBasics s conn
                          stmtExecute stmt ps
                    , stmtReset = stmtReset stmt
                    , stmtFinalize = stmtFinalize stmt
                    }
            }
  
        conn' =
          conn
            { connHooks = hooks
            , connBegin = \f mIso -> do
                ctxt <- getContext
                s <- createSpan t ctxt "transaction" (defaultSpanArguments {kind = Client, attributes = attrs})
                annotateBasics s conn
                writeIORef connSpanInFlight (Just s)
                writeIORef connParentSpan (lookupSpan ctxt)
                adjustContext (insertSpan s)
                case mIso of
                  Nothing -> pure ()
                  Just iso -> addAttribute s "db.transaction.isolation" $ case iso of
                    ReadUncommitted -> "read uncommitted" :: Text
                    ReadCommitted -> "read committed"
                    RepeatableRead -> "repeatable read"
                    Serializable -> "serializable"
                connBegin conn f mIso
            , connCommit = \f -> do
                spanInFlight <- readIORef connSpanInFlight
                parentSpan <- readIORef connParentSpan
                let act = do
                      (Timestamp tsStart) <- getTimestamp
                      result <- tryAny $ connCommit conn f
                      (Timestamp tsEnd) <- getTimestamp
                      forM_ spanInFlight $ \s -> do
                        addAttributes
                          s
                          [ ("db.transaction.outcome", toAttribute ("committed" :: Text))
                          , ("db.transaction.commit_duration_ns", toAttribute $ fromIntegral @Integer @Int $ toNanoSecs (diffTimeSpec tsStart tsEnd) `div` 1000)
                          ]
                        endSpan s Nothing
                        case result of
                          Left (SomeException err) -> do
                            recordException s [("exception.escaped", toAttribute True)] Nothing err
                            throwIO err
                          Right _ -> pure ()
                act `finally` do
                  adjustContext $ \ctx ->
                    maybe (removeSpan ctx) (`insertSpan` ctx) parentSpan
                  forM_ spanInFlight $ \s -> endSpan s Nothing
            , connRollback = \f -> do
                spanInFlight <- readIORef connSpanInFlight
                parentSpan <- readIORef connParentSpan
                let act = do
                      (Timestamp tsStart) <- getTimestamp
                      result <- tryAny $ connRollback conn f
                      e@(Timestamp tsEnd) <- getTimestamp
                      forM_ spanInFlight $ \s -> do
                        addAttributes
                          s
                          [ ("db.transaction.outcome", toAttribute ("rolled back" :: Text))
                          , ("db.transaction.commit_duration_microseconds", toAttribute $ fromIntegral @Integer @Int $ toNanoSecs (diffTimeSpec tsStart tsEnd `div` 1000))
                          ]
                        endSpan s (Just e)
                        case result of
                          Left (SomeException err) -> do
                            recordException s [("exception.escaped", toAttribute True)] Nothing err
                            throwIO err
                          Right _ -> pure ()
                act `finally` do
                  adjustContext $ \ctx ->
                    maybe (removeSpan ctx) (`insertSpan` ctx) parentSpan
                  forM_ spanInFlight $ \s -> endSpan s Nothing
            , -- TODO: This doesn't work when we wrap the connections for the pool.
              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)
      ]
  