-- | PostgreSQL logging back-end.
module Log.Backend.PostgreSQL (withPgLogger) where

import Control.Concurrent
import Control.Exception.Lifted
import Control.Monad
import Control.Monad.State.Lazy
import Control.Monad.IO.Unlift
import Data.Aeson ((.=), Value(..), object, encode)
import Data.List.Split
import Data.Monoid.Utils
import Data.String
import Data.Typeable
import Database.PostgreSQL.PQTypes
import qualified Data.ByteString.Base64 as B64
import qualified Data.Foldable as Foldable
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Vector as V

import Log.Data
import Log.Logger
import Log.Internal.Logger
import qualified Log.Internal.Aeson.Compat as AC

newtype InvalidEncodingRecoveryAttempt = Attempt Int
  deriving Int -> InvalidEncodingRecoveryAttempt
InvalidEncodingRecoveryAttempt -> Int
InvalidEncodingRecoveryAttempt -> [InvalidEncodingRecoveryAttempt]
InvalidEncodingRecoveryAttempt -> InvalidEncodingRecoveryAttempt
InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
$cenumFromThenTo :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
enumFromTo :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
$cenumFromTo :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
enumFromThen :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
$cenumFromThen :: InvalidEncodingRecoveryAttempt
-> InvalidEncodingRecoveryAttempt
-> [InvalidEncodingRecoveryAttempt]
enumFrom :: InvalidEncodingRecoveryAttempt -> [InvalidEncodingRecoveryAttempt]
$cenumFrom :: InvalidEncodingRecoveryAttempt -> [InvalidEncodingRecoveryAttempt]
fromEnum :: InvalidEncodingRecoveryAttempt -> Int
$cfromEnum :: InvalidEncodingRecoveryAttempt -> Int
toEnum :: Int -> InvalidEncodingRecoveryAttempt
$ctoEnum :: Int -> InvalidEncodingRecoveryAttempt
pred :: InvalidEncodingRecoveryAttempt -> InvalidEncodingRecoveryAttempt
$cpred :: InvalidEncodingRecoveryAttempt -> InvalidEncodingRecoveryAttempt
succ :: InvalidEncodingRecoveryAttempt -> InvalidEncodingRecoveryAttempt
$csucc :: InvalidEncodingRecoveryAttempt -> InvalidEncodingRecoveryAttempt
Enum

-- | Create a 'pgLogger' for the duration of the given action, and
-- shut it down afterwards, making sure that all buffered messages are
-- actually written to the DB.
withPgLogger :: MonadUnliftIO m => ConnectionSourceM IO -> (Logger -> m r) -> m r
withPgLogger :: forall (m :: * -> *) r.
MonadUnliftIO m =>
ConnectionSourceM IO -> (Logger -> m r) -> m r
withPgLogger ConnectionSourceM IO
cs Logger -> m r
act = forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
unlift -> do
  Logger
logger <- ConnectionSourceM IO -> IO Logger
pgLogger ConnectionSourceM IO
cs
  forall r. Logger -> (Logger -> IO r) -> IO r
withLogger Logger
logger (forall a. m a -> IO a
unlift forall b c a. (b -> c) -> (a -> b) -> a -> c
. Logger -> m r
act)

-- | Start an asynchronous logger thread that inserts log messages
-- into a PostgreSQL database.
--
-- Please use 'withPglogger' instead, which is more exception-safe
-- (see the note attached to 'mkBulkLogger').
pgLogger :: ConnectionSourceM IO -> IO Logger
pgLogger :: ConnectionSourceM IO -> IO Logger
pgLogger ConnectionSourceM IO
cs = Text -> ([LogMessage] -> IO ()) -> IO () -> IO Logger
mkBulkLogger forall s. IsString s => s
loggerName
              (forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
serialize forall a b. (a -> b) -> a -> b
$ Int -> InvalidEncodingRecoveryAttempt
Attempt Int
1) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e. Int -> [e] -> [[e]]
chunksOf Int
1000)
              (forall (m :: * -> *) a. Monad m => a -> m a
return ())
  where
    loggerName :: IsString s => s
    loggerName :: forall s. IsString s => s
loggerName = s
"PostgreSQL"

    sqlInsertLog :: SQL
    sqlInsertLog :: SQL
sqlInsertLog = SQL
"INSERT INTO logs "
      forall m. (IsString m, Monoid m) => m -> m -> m
<+> SQL
"(time, level, component, domain, message, data) VALUES"

    serialize :: InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
    serialize :: InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
serialize !InvalidEncodingRecoveryAttempt
attempt [LogMessage]
msgs = forall (m :: * -> *) a.
(MonadBase IO m, MonadMask m) =>
ConnectionSourceM m -> TransactionSettings -> DBT m a -> m a
runDBT ConnectionSourceM IO
cs TransactionSettings
ts
      (forall (m :: * -> *). MonadDB m => SQL -> m ()
runSQL_ forall a b. (a -> b) -> a -> b
$ SQL
sqlInsertLog
       forall m. (IsString m, Monoid m) => m -> m -> m
<+> forall m. Monoid m => m -> [m] -> m
mintercalate SQL
", " (forall a b. (a -> b) -> [a] -> [b]
map LogMessage -> SQL
sqlifyMessage [LogMessage]
msgs))
      forall (m :: * -> *) a.
MonadBaseControl IO m =>
m a -> [Handler m a] -> m a
`catches` [
        -- Propagate base async exceptions thrown by the runtime system.
        forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler forall a b. (a -> b) -> a -> b
$ \(AsyncException
e::AsyncException) -> forall (m :: * -> *) e a. (MonadBase IO m, Exception e) => e -> m a
throwIO AsyncException
e
      , forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler forall a b. (a -> b) -> a -> b
$ \(SomeException
e::SomeException) -> case forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e of
        Just dbe :: DBException
dbe@DBException{}
          | Just DetailedQueryError
qe <- DBException -> Maybe DetailedQueryError
getEncodingQueryError DBException
dbe -> case InvalidEncodingRecoveryAttempt
attempt of
            Attempt Int
1 -> do
              -- If a client uses UTF-8 encoding (TODO: in fact it should
              -- always be the case as Text is encoded as UTF-8 for sql
              -- serialization), then this error occurs only when any of the
              -- strings we want to serialize contains NULL bytes. In such
              -- case we scan the logs and replace each NULL with "\0".
              [Char] -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall s. IsString s => s
loggerName
                forall a. [a] -> [a] -> [a]
++ [Char]
": couldn't serialize logs due to character encoding error \""
                forall a. [a] -> [a] -> [a]
++ DetailedQueryError -> [Char]
qeMessagePrimary DetailedQueryError
qe forall a. [a] -> [a] -> [a]
++ [Char]
"\", removing NULL bytes and retrying"
              InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
serialize (forall a. Enum a => a -> a
succ InvalidEncodingRecoveryAttempt
attempt) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\LogMessage
msg ->
                -- If any text inside the message had NULL bytes,
                -- add acknowledgment of that fact to its data.
                case forall s a. State s a -> s -> (a, s)
runState (forall (m :: * -> *).
(Applicative m, Monad m) =>
(Text -> m Text) -> LogMessage -> m LogMessage
mapTexts Text -> State Bool Text
removeNULLs LogMessage
msg) Bool
False of
                  (LogMessage
newMsg, Bool
True) -> LogMessage
newMsg {
                    lmData :: Value
lmData = LogMessage -> Value
lmData LogMessage
newMsg
                             Value -> (Text, Value) -> Value
`addPair` (Text
"_log", Value
"NULL bytes were escaped")
                  }
                  (LogMessage
_, Bool
False) -> LogMessage
msg) [LogMessage]
msgs
            Attempt Int
2 -> do
              -- This should never happen, but let us be paranoid for
              -- a minute. If the error is still happening after removal
              -- of NULL bytes, go through each message and encode all
              -- texts as base64, effectively transforming them into ASCII.
              [Char] -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall s. IsString s => s
loggerName
                forall a. [a] -> [a] -> [a]
++ [Char]
": couldn't serialize logs due to character encoding error \""
                forall a. [a] -> [a] -> [a]
++ DetailedQueryError -> [Char]
qeMessagePrimary DetailedQueryError
qe
                forall a. [a] -> [a] -> [a]
++ [Char]
"\" after NULL bytes were removed, encoding all texts"
                forall a. [a] -> [a] -> [a]
++ [Char]
" in the problematic batch as base64 to make them ASCII"
              InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
serialize (forall a. Enum a => a -> a
succ InvalidEncodingRecoveryAttempt
attempt) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\LogMessage
msg ->
                let newMsg :: LogMessage
newMsg = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
(Applicative m, Monad m) =>
(Text -> m Text) -> LogMessage -> m LogMessage
mapTexts Text -> Identity Text
convertBase64 LogMessage
msg
                 in LogMessage
newMsg {
                  lmData :: Value
lmData = LogMessage -> Value
lmData LogMessage
newMsg
                    Value -> (Text, Value) -> Value
`addPair` (Text
"_log", Value
"Texts encoded as base64")
                }) [LogMessage]
msgs
            Attempt Int
_ -> do
              -- This can't happen, all texts are ASCII now.
              [Char] -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall s. IsString s => s
loggerName
                forall a. [a] -> [a] -> [a]
++ [Char]
": impossible happened "
                forall a. [a] -> [a] -> [a]
++ [Char]
"(>2 attempt failed because of character encoding error \""
                forall a. [a] -> [a] -> [a]
++ DetailedQueryError -> [Char]
qeMessagePrimary DetailedQueryError
qe
                forall a. [a] -> [a] -> [a]
++ [Char]
"\" even though all texts are ASCII), skipping the batch"
        Maybe DBException
_ -> do
          [Char] -> IO ()
putStrLn forall a b. (a -> b) -> a -> b
$ forall s. IsString s => s
loggerName
            forall a. [a] -> [a] -> [a]
++ [Char]
": couldn't serialize logs:"
            forall m. (IsString m, Monoid m) => m -> m -> m
<+> forall a. Show a => a -> [Char]
show SomeException
e forall a. [a] -> [a] -> [a]
++ [Char]
", retrying in 10 seconds"
          Int -> IO ()
threadDelay forall a b. (a -> b) -> a -> b
$ Int
10 forall a. Num a => a -> a -> a
* Int
1000000
          -- Do not increment the attempt here, it's used to
          -- track invalid encoding recovery attempts only.
          InvalidEncodingRecoveryAttempt -> [LogMessage] -> IO ()
serialize InvalidEncodingRecoveryAttempt
attempt [LogMessage]
msgs
      ]

    addPair :: Value -> (T.Text, Value) -> Value
    addPair :: Value -> (Text, Value) -> Value
addPair Value
data_ (Text
name, Value
value) = case Value
data_ of
      Object Object
obj -> Object -> Value
Object forall a b. (a -> b) -> a -> b
$ forall v. Key -> v -> KeyMap v -> KeyMap v
AC.insert (Text -> Key
AC.fromText Text
name) Value
value Object
obj
      Value
_          -> [Pair] -> Value
object
                    [ Key
"_data" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Value
data_
                    , Key
"_log"  forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Value
value
                    ]

    getEncodingQueryError :: DBException -> Maybe DetailedQueryError
    getEncodingQueryError :: DBException -> Maybe DetailedQueryError
getEncodingQueryError DBException{e
sql
dbeQueryContext :: ()
dbeError :: ()
dbeError :: e
dbeQueryContext :: sql
..}
      | Just (DetailedQueryError
qe::DetailedQueryError) <- forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast e
dbeError
      ,    DetailedQueryError -> ErrorCode
qeErrorCode DetailedQueryError
qe forall a. Eq a => a -> a -> Bool
== ErrorCode
CharacterNotInRepertoire
        Bool -> Bool -> Bool
|| DetailedQueryError -> ErrorCode
qeErrorCode DetailedQueryError
qe forall a. Eq a => a -> a -> Bool
== ErrorCode
UntranslatableCharacter = forall a. a -> Maybe a
Just DetailedQueryError
qe
      | Bool
otherwise = forall a. Maybe a
Nothing

    convertBase64 :: T.Text -> Identity T.Text
    convertBase64 :: Text -> Identity Text
convertBase64 = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
T.decodeLatin1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
B64.encode forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf8

    removeNULLs :: T.Text -> State Bool T.Text
    removeNULLs :: Text -> State Bool Text
removeNULLs Text
s = do
      let newS :: Text
newS = Text -> Text -> Text -> Text
T.replace Text
"\0" Text
"\\0" Text
s
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Text -> Int
T.length Text
newS forall a. Eq a => a -> a -> Bool
/= Text -> Int
T.length Text
s) forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *). MonadState s m => s -> m ()
put Bool
True
      forall (m :: * -> *) a. Monad m => a -> m a
return Text
newS

    mapTexts :: forall m. (Applicative m, Monad m)
             => (T.Text -> m T.Text) -> LogMessage -> m LogMessage
    mapTexts :: forall (m :: * -> *).
(Applicative m, Monad m) =>
(Text -> m Text) -> LogMessage -> m LogMessage
mapTexts Text -> m Text
doText LogMessage
lm = do
      Text
component <- Text -> m Text
doText      forall a b. (a -> b) -> a -> b
$ LogMessage -> Text
lmComponent LogMessage
lm
      [Text]
domain    <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Text -> m Text
doText forall a b. (a -> b) -> a -> b
$ LogMessage -> [Text]
lmDomain LogMessage
lm
      Text
message   <- Text -> m Text
doText      forall a b. (a -> b) -> a -> b
$ LogMessage -> Text
lmMessage LogMessage
lm
      Value
data_     <- Value -> m Value
doValue     forall a b. (a -> b) -> a -> b
$ LogMessage -> Value
lmData LogMessage
lm
      forall (m :: * -> *) a. Monad m => a -> m a
return LogMessage
lm {
        lmComponent :: Text
lmComponent = Text
component
      , lmDomain :: [Text]
lmDomain    = [Text]
domain
      , lmMessage :: Text
lmMessage   = Text
message
      , lmData :: Value
lmData      = Value
data_
      }
      where
        doValue :: Value -> m Value
        doValue :: Value -> m Value
doValue (Object Object
obj) =
          Object -> Value
Object forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
Foldable.foldrM
          (\(Key
name, Value
value) Object
acc ->
             forall v. Key -> v -> KeyMap v -> KeyMap v
AC.insert forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). Monad m => (Text -> m Text) -> Key -> m Key
AC.doName Text -> m Text
doText Key
name forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Value -> m Value
doValue Value
value forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Object
acc)
          forall v. KeyMap v
AC.empty (forall v. KeyMap v -> [(Key, v)]
AC.toList Object
obj)
        doValue (Array Array
arr)  = Array -> Value
Array forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM Value -> m Value
doValue Array
arr
        doValue (String Text
s)   = Text -> Value
String forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> m Text
doText Text
s
        doValue Value
v            = forall (m :: * -> *) a. Monad m => a -> m a
return Value
v

    sqlifyMessage :: LogMessage -> SQL
    sqlifyMessage :: LogMessage -> SQL
sqlifyMessage LogMessage{[Text]
UTCTime
Text
Value
LogLevel
lmTime :: LogMessage -> UTCTime
lmLevel :: LogMessage -> LogLevel
lmData :: Value
lmMessage :: Text
lmLevel :: LogLevel
lmTime :: UTCTime
lmDomain :: [Text]
lmComponent :: Text
lmMessage :: LogMessage -> Text
lmDomain :: LogMessage -> [Text]
lmComponent :: LogMessage -> Text
lmData :: LogMessage -> Value
..} = forall a. Monoid a => [a] -> a
mconcat [
        SQL
"("
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> UTCTime
lmTime
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> LogLevel -> Text
showLogLevel LogLevel
lmLevel
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> Text
lmComponent
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> forall a. [a] -> Array1 a
Array1 [Text]
lmDomain
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> Text
lmMessage
      , SQL
"," forall t. (Show t, ToSQL t) => SQL -> t -> SQL
<?> forall jsonb. jsonb -> JSONB jsonb
JSONB (forall a. ToJSON a => a -> ByteString
encode Value
lmData)
      , SQL
")"
      ]

    ts :: TransactionSettings
    ts :: TransactionSettings
ts = TransactionSettings
defaultTransactionSettings {
      tsAutoTransaction :: Bool
tsAutoTransaction = Bool
False
    }