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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module AppT where

import Builder
import Control.Monad.Logger qualified as Logger
import Control.Monad.Logger.CallStack
import Control.Monad.Reader
import Data.Aeson qualified as Json
import Data.Error.Tree
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as HashMap
import Data.Pool (Pool)
import Data.String (IsString (fromString))
import Data.Text qualified as Text
import Data.Time (UTCTime)
import Database.PostgreSQL.Simple qualified as Postgres
import FieldParser (FieldParser)
import FieldParser qualified as Field
import GHC.Generics qualified as G
import GHC.Records (getField)
import GHC.Stack qualified
import GHC.TypeLits
import Json.Enc
import Json.Enc qualified as Enc
import Label
import MyPrelude
import OpenTelemetry.Context.ThreadLocal qualified as Otel
import OpenTelemetry.Trace qualified as Otel hiding (getTracer, inSpan, inSpan')
import OpenTelemetry.Trace.Core qualified as Otel hiding (inSpan, inSpan')
import OpenTelemetry.Trace.Monad qualified as Otel
import Postgres.MonadPostgres
import Pretty qualified
import System.IO qualified as IO
import Tool (Tool)
import UnliftIO
import Prelude hiding (span)

data Context = Context
  { pgConfig ::
      T2
        "logDatabaseQueries"
        DebugLogDatabaseQueries
        "prettyPrintDatabaseQueries"
        PrettyPrintDatabaseQueries,
    pgConnPool :: (Pool Postgres.Connection),
    tracer :: Otel.Tracer,
    transmissionSessionId :: IORef (Maybe ByteString),
    redactedApiKey :: ByteString,
    tools :: Tools,
    transmissionDownloads :: Maybe (T2 "downloadDirectory" FilePath "staticFileEndpoint" Text)
  }

newtype AppT m a = AppT {unAppT :: ReaderT Context m a}
  deriving newtype (Functor, Applicative, Monad, MonadIO, MonadUnliftIO, MonadThrow)

type App a = AppT IO a

type AppTransaction a = Transaction (AppT IO) a

data AppException
  = AppExceptionTree ErrorTree
  | AppExceptionPretty [Pretty.Err]
  | AppExceptionEnc Enc
  deriving anyclass (Exception)

instance IsString AppException where
  fromString s = AppExceptionTree (fromString s)

instance Show AppException where
  showsPrec _ (AppExceptionTree t) = ("AppException: " ++) . ((textToString $ prettyErrorTree t) ++)
  showsPrec _ (AppExceptionPretty t) = ("AppException: " ++) . ((Pretty.prettyErrsNoColor t) ++)
  showsPrec _ (AppExceptionEnc e) = ((textToString $ Enc.encToTextPretty e) ++)

instance (MonadIO m) => MonadLogger (AppT m) where
  monadLoggerLog loc src lvl msg = liftIO $ Logger.defaultOutput IO.stderr loc src lvl (Logger.toLogStr msg)

instance (Monad m) => Otel.MonadTracer (AppT m) where
  getTracer = AppT $ asks (.tracer)

data Tools = Tools
  { exiftool :: Tool
  }

getTools :: (Monad m) => AppT m Tools
getTools = AppT $ asks (.tools)

class (MonadUnliftIO m, Otel.MonadTracer m) => MonadOtel m

instance (MonadUnliftIO m) => MonadOtel (AppT m)

instance (MonadOtel m) => MonadOtel (Transaction m)

inSpan :: (MonadOtel m) => Text -> m a -> m a
inSpan name = Otel.inSpan name Otel.defaultSpanArguments

inSpan' :: (MonadOtel m) => Text -> (Otel.Span -> m a) -> m a
inSpan' name = Otel.inSpan' name Otel.defaultSpanArguments

-- | Record an exception in the current span, and set the span status to error
recordError :: (MonadOtel m, MonadLogger m) => Otel.Span -> Error -> m ()
recordError span err = do
  logError $ prettyError err
  Otel.addEvent span $
    Otel.NewEvent
      { newEventName = "exception",
        newEventAttributes =
          HashMap.fromList
            [ ("exception.message", Otel.toAttribute $ prettyError err)
            -- ("exception.stacktrace", A.toAttribute $ T.unlines $ map T.pack cs)
            ],
        newEventTimestamp = Nothing
      }
  Otel.setStatus
    span
    -- dunno what the difference is between setting exception.message and this string,
    -- but if we leave it empty here `otel.status_description` is not set.
    (Otel.Error "")

-- | Record an exception in the current span, and set the span status to error
recordErrorTree :: (MonadOtel m, MonadLogger m) => Otel.Span -> ErrorTree -> m ()
recordErrorTree span e@(ErrorTree errTree) = do
  logError $ prettyErrorTree e
  Otel.addEvent span $
    Otel.NewEvent
      { newEventName = "exception",
        newEventAttributes =
          HashMap.fromList
            [ ( "exception.message",
                toOtelJsonAttr $
                  Enc.encoding (Json.toEncoding (errTree <&> prettyError))
              )
              -- ("exception.stacktrace", A.toAttribute $ T.unlines $ map T.pack cs)
            ],
        newEventTimestamp = Nothing
      }
  Otel.setStatus
    span
    -- dunno what the difference is between setting exception.message and this string,
    -- but if we leave it empty here `otel.status_description` is not set.
    (Otel.Error "")

-- | Add the attribute to the span, prefixing it with the `_` namespace (to easier distinguish our application’s tags from standard tags)
addAttribute :: (MonadIO m, Otel.ToAttribute a) => Otel.Span -> Text -> a -> m ()
addAttribute span key a = Otel.addAttribute span ("_." <> key) a

-- | Add the attributes to the span, prefixing each key with the `_` namespace (to easier distinguish our application’s tags from standard tags)
addAttributes :: (MonadIO m) => Otel.Span -> HashMap Text Otel.Attribute -> m ()
addAttributes span attrs = Otel.addAttributes span $ attrs & HashMap.mapKeys ("_." <>)

addEventSimple :: (MonadIO m) => Otel.Span -> Text -> m ()
addEventSimple span name =
  Otel.addEvent
    span
    Otel.NewEvent
      { Otel.newEventName = name,
        Otel.newEventTimestamp = Nothing,
        Otel.newEventAttributes = mempty
      }

-- | Create an otel attribute from a json encoder
jsonAttribute :: Enc -> Otel.Attribute
jsonAttribute e = e & Enc.encToTextPretty & Otel.toAttribute

instance Otel.ToAttribute (a, Build Txt a) where
  toAttribute (a, b) = buildText b a & Otel.toAttribute

parseOrThrow :: (MonadThrow m, MonadIO m) => Otel.Span -> FieldParser from to -> from -> m to
parseOrThrow span fp f =
  f & Field.runFieldParser fp & \case
    Left err -> appThrow span (AppExceptionTree $ singleError err)
    Right a -> pure a

orThrowAppErrorNewSpan :: (MonadThrow m, MonadOtel m) => Text -> Either AppException a -> m a
orThrowAppErrorNewSpan msg = \case
  Left err -> appThrowNewSpan msg err
  Right a -> pure a

appThrowNewSpan :: (MonadThrow m, MonadOtel m) => Text -> AppException -> m a
appThrowNewSpan spanName exc = inSpan' spanName $ \span -> do
  let msg = case exc of
        AppExceptionTree e -> prettyErrorTree e
        AppExceptionPretty p -> Pretty.prettyErrsNoColor p & stringToText
        AppExceptionEnc e -> Enc.encToTextPretty e
  recordException
    span
    ( T2
        (label @"type_" "AppException")
        (label @"message" msg)
    )
  throwM $ exc

appThrow :: (MonadThrow m, MonadIO m) => Otel.Span -> AppException -> m a
appThrow span exc = do
  let msg = case exc of
        AppExceptionTree e -> prettyErrorTree e
        AppExceptionPretty p -> Pretty.prettyErrsNoColor p & stringToText
        AppExceptionEnc e -> Enc.encToTextPretty e
  recordException
    span
    ( T2
        (label @"type_" "AppException")
        (label @"message" msg)
    )
  throwM $ exc

orAppThrow :: (MonadThrow m, MonadIO m) => Otel.Span -> Either AppException a -> m a
orAppThrow span = \case
  Left err -> appThrow span err
  Right a -> pure a

-- | If action returns a Left, throw an AppException
assertM :: (MonadThrow f, MonadIO f) => Otel.Span -> (t -> Either AppException a) -> t -> f a
assertM span f v = case f v of
  Right a -> pure a
  Left err -> appThrow span err

assertMNewSpan :: (MonadThrow f, MonadOtel f) => Text -> (t -> Either AppException a) -> t -> f a
assertMNewSpan spanName f v = case f v of
  Right a -> pure a
  Left err -> appThrowNewSpan spanName err

-- | A specialized variant of @addEvent@ that records attributes conforming to
-- the OpenTelemetry specification's
-- <https://github.com/open-telemetry/opentelemetry-specification/blob/49c2f56f3c0468ceb2b69518bcadadd96e0a5a8b/specification/trace/semantic_conventions/exceptions.md semantic conventions>
--
-- @since 0.0.1.0
recordException ::
  ( MonadIO m,
    HasField "message" r Text,
    HasField "type_" r Text
  ) =>
  Otel.Span ->
  r ->
  m ()
recordException span dat = liftIO $ do
  callStack <- GHC.Stack.whoCreated dat.message
  newEventTimestamp <- Just <$> Otel.getTimestamp
  Otel.addEvent span $
    Otel.NewEvent
      { newEventName = "exception",
        newEventAttributes =
          HashMap.fromList
            [ ("exception.type", Otel.toAttribute @Text dat.type_),
              ("exception.message", Otel.toAttribute @Text dat.message),
              ("exception.stacktrace", Otel.toAttribute @Text $ Text.unlines $ Prelude.map stringToText callStack)
            ],
        ..
      }

-- * Async wrappers with Otel tracing

withAsyncTraced :: (MonadUnliftIO m) => m a -> (Async a -> m b) -> m b
withAsyncTraced act f = do
  ctx <- Otel.getContext
  withAsync
    ( do
        _old <- Otel.attachContext ctx
        act
    )
    f

-- | Run two actions concurrently, and add them to the current Otel trace
concurrentlyTraced :: (MonadUnliftIO m) => m a -> m b -> m (a, b)
concurrentlyTraced act1 act2 = do
  ctx <- Otel.getContext
  concurrently
    ( do
        _old <- Otel.attachContext ctx
        act1
    )
    ( do
        _old <- Otel.attachContext ctx
        act2
    )

mapConcurrentlyTraced :: (MonadUnliftIO m, Traversable t) => (a -> m b) -> t a -> m (t b)
mapConcurrentlyTraced f t = do
  ctx <- Otel.getContext
  mapConcurrently
    ( \a -> do
        _old <- Otel.attachContext ctx
        f a
    )
    t

-- * Postgres

instance (MonadThrow m, MonadUnliftIO m) => MonadPostgres (AppT m) where
  execute = executeImpl dbConfig
  executeMany = executeManyImpl dbConfig
  executeManyReturningWith = executeManyReturningWithImpl dbConfig
  queryWith = queryWithImpl dbConfig
  queryWith_ = queryWithImpl_ (dbConfig <&> snd)

  foldRowsWithAcc = foldRowsWithAccImpl dbConfig
  runTransaction = runPGTransaction

dbConfig :: (Monad m) => AppT m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries)
dbConfig =
  AppT $
    asks
      ( \c ->
          ( c.pgConfig.logDatabaseQueries,
            c.pgConfig.prettyPrintDatabaseQueries
          )
      )

runPGTransaction :: (MonadUnliftIO m) => Transaction (AppT m) a -> AppT m a
runPGTransaction (Transaction transaction) = do
  pool <- AppT ask <&> (.pgConnPool)
  withRunInIO $ \unliftIO ->
    withPGTransaction pool $ \conn -> do
      unliftIO $ runReaderT transaction conn

-- | Best effort to convert a value to a JSON string that can be put in an Otel attribute.
toOtelJsonAttr :: (ToOtelJsonAttr a) => a -> Otel.Attribute
toOtelJsonAttr = toOtelJsonAttrImpl >>> Enc.encToTextPretty >>> Otel.toAttribute

-- | Best effort to convert a value to a JSON string that can be put in an Otel attribute.
class ToOtelJsonAttr a where
  toOtelJsonAttrImpl :: a -> Enc

instance ToOtelJsonAttr Enc where
  toOtelJsonAttrImpl = id

-- | Bytes are leniently converted to Text, because they are often used as UTF-8 encoded strings.
instance ToOtelJsonAttr ByteString where
  toOtelJsonAttrImpl = Enc.text . bytesToTextUtf8Lenient

instance ToOtelJsonAttr Text where
  toOtelJsonAttrImpl = Enc.text

instance ToOtelJsonAttr Int where
  toOtelJsonAttrImpl = Enc.int

instance ToOtelJsonAttr Natural where
  toOtelJsonAttrImpl = Enc.natural

instance ToOtelJsonAttr Bool where
  toOtelJsonAttrImpl = Enc.bool

instance ToOtelJsonAttr UTCTime where
  toOtelJsonAttrImpl = Enc.text . showToText

instance (ToOtelJsonAttr a) => ToOtelJsonAttr (Maybe a) where
  toOtelJsonAttrImpl = \case
    Nothing -> Enc.null
    Just a -> toOtelJsonAttrImpl a

instance (ToOtelJsonAttr a) => ToOtelJsonAttr [a] where
  toOtelJsonAttrImpl = Enc.list toOtelJsonAttrImpl

instance (ToOtelJsonAttr t1, ToOtelJsonAttr t2, KnownSymbol l1, KnownSymbol l2) => ToOtelJsonAttr (T2 l1 t1 l2 t2) where
  toOtelJsonAttrImpl (T2 a b) =
    Enc.object
      [ (symbolText @l1, a & getField @l1 & toOtelJsonAttrImpl),
        (symbolText @l2, b & getField @l2 & toOtelJsonAttrImpl)
      ]

instance (ToOtelJsonAttr t1, ToOtelJsonAttr t2, ToOtelJsonAttr t3, KnownSymbol l1, KnownSymbol l2, KnownSymbol l3) => ToOtelJsonAttr (T3 l1 t1 l2 t2 l3 t3) where
  toOtelJsonAttrImpl (T3 a b c) =
    Enc.object
      [ (symbolText @l1, a & getField @l1 & toOtelJsonAttrImpl),
        (symbolText @l2, b & getField @l2 & toOtelJsonAttrImpl),
        (symbolText @l3, c & getField @l3 & toOtelJsonAttrImpl)
      ]

instance (ToOtelJsonAttr t1, ToOtelJsonAttr t2) => ToOtelJsonAttr (t1, t2) where
  toOtelJsonAttrImpl t = Enc.tuple2 toOtelJsonAttrImpl toOtelJsonAttrImpl t

instance (ToOtelJsonAttr t1, ToOtelJsonAttr t2, ToOtelJsonAttr t3) => ToOtelJsonAttr (t1, t2, t3) where
  toOtelJsonAttrImpl t = Enc.tuple3 toOtelJsonAttrImpl toOtelJsonAttrImpl toOtelJsonAttrImpl t

-- | Pretty-print the given value to a string
toOtelAttrGenericStruct :: (Generic a, GenericStructSimple (G.Rep a)) => a -> Otel.Attribute
toOtelAttrGenericStruct a = toOtelJsonAttr @Enc $ encodeSimpleValue $ G.from a

class GenericStruct f where
  encodeStructAsObject :: f a -> [(Text, Enc)]

-- :*: (product)
-- Object fields (get field name and put into a list of key-value pair)
instance
  (KnownSymbol l, ToOtelJsonAttr val) =>
  GenericStruct (G.M1 G.S (G.MetaSel (Just l) u s f) (G.K1 i val))
  where
  encodeStructAsObject (G.M1 (G.K1 x)) = [(symbolText @l, toOtelJsonAttrImpl x)]

-- Concatenate two fields in a struct
instance (GenericStruct f, GenericStruct g) => GenericStruct (f G.:*: g) where
  encodeStructAsObject (f G.:*: g) = encodeStructAsObject f <> encodeStructAsObject g

class GenericStructSimple f where
  encodeSimpleValue :: f a -> Enc

instance
  (ToOtelJsonAttr val, KnownSymbol l) =>
  GenericStructSimple (G.M1 G.S (G.MetaSel (Just l) u s f) (G.K1 i val))
  where
  encodeSimpleValue (G.M1 x) = Enc.object $ [(symbolText @l, encodeSimpleValue x)]

-- pass through other M1
instance (GenericStructSimple f) => GenericStructSimple (G.M1 G.D u f) where
  encodeSimpleValue (G.M1 x) = encodeSimpleValue x

-- pass through other M1
instance (GenericStructSimple f) => GenericStructSimple (G.M1 G.C u f) where
  encodeSimpleValue (G.M1 x) = encodeSimpleValue x

-- | Encode a generic representation as an object with :*:
instance (GenericStruct f, GenericStruct g) => GenericStructSimple (f G.:*: g) where
  encodeSimpleValue (a G.:*: b) = Enc.object $ encodeStructAsObject a <> encodeStructAsObject b

-- Void
instance GenericStructSimple G.V1 where
  encodeSimpleValue x = case x of {}

-- Empty type is the empty object
instance GenericStructSimple G.U1 where
  encodeSimpleValue _ = emptyObject

-- K1
instance (ToOtelJsonAttr val) => GenericStructSimple (G.K1 i val) where
  encodeSimpleValue (G.K1 x) = toOtelJsonAttrImpl x