{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NamedFieldPuns #-}

{- |
Module      : Database.DuckDB.Simple.Logging
Description : High-level wrappers for DuckDB custom log storage.
-}
module Database.DuckDB.Simple.Logging (
    LogEntry (..),
    registerLogStorage,
) where

import Control.Exception (bracket, onException)
import Data.Ratio ((%))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Foreign as TextForeign
import Data.Time.Clock (UTCTime)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
import Database.DuckDB.FFI
import Database.DuckDB.Simple.Internal (Connection, mkDeleteCallback, throwRegistrationError, withDatabaseHandle)
import Foreign.C.String (CString, peekCString)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr, freeHaskellFunPtr, nullPtr)
import Foreign.StablePtr (StablePtr, castPtrToStablePtr, castStablePtrToPtr, deRefStablePtr, freeStablePtr, newStablePtr)
import Foreign.Storable (peek, poke)

-- | A single log event delivered through DuckDB's log-storage callback.
data LogEntry = LogEntry
    { LogEntry -> Maybe UTCTime
logEntryTimestamp :: !(Maybe UTCTime)
    , LogEntry -> Text
logEntryLevel :: !Text
    , LogEntry -> Text
logEntryType :: !Text
    , LogEntry -> Text
logEntryMessage :: !Text
    }
    deriving (LogEntry -> LogEntry -> Bool
(LogEntry -> LogEntry -> Bool)
-> (LogEntry -> LogEntry -> Bool) -> Eq LogEntry
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LogEntry -> LogEntry -> Bool
== :: LogEntry -> LogEntry -> Bool
$c/= :: LogEntry -> LogEntry -> Bool
/= :: LogEntry -> LogEntry -> Bool
Eq, Int -> LogEntry -> ShowS
[LogEntry] -> ShowS
LogEntry -> String
(Int -> LogEntry -> ShowS)
-> (LogEntry -> String) -> ([LogEntry] -> ShowS) -> Show LogEntry
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LogEntry -> ShowS
showsPrec :: Int -> LogEntry -> ShowS
$cshow :: LogEntry -> String
show :: LogEntry -> String
$cshowList :: [LogEntry] -> ShowS
showList :: [LogEntry] -> ShowS
Show)

-- | Register a custom log storage callback on the database behind a connection.
registerLogStorage :: Connection -> Text -> (LogEntry -> IO ()) -> IO ()
registerLogStorage :: Connection -> Text -> (LogEntry -> IO ()) -> IO ()
registerLogStorage Connection
conn Text
name LogEntry -> IO ()
callback = do
    writeCb <- (Ptr ()
 -> Ptr DuckDBTimestamp -> CString -> CString -> CString -> IO ())
-> IO DuckDBLoggerWriteLogEntryFun
mkWriteLogEntryCallback ((LogEntry -> IO ())
-> Ptr ()
-> Ptr DuckDBTimestamp
-> CString
-> CString
-> CString
-> IO ()
logStorageHandler LogEntry -> IO ()
callback)
    callbackStable <- newStablePtr writeCb
    deleteCb <- mkDeleteCallback releaseWriteLogCallback
    let release = DuckDBLoggerWriteLogEntryFun -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBLoggerWriteLogEntryFun
writeCb IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> StablePtr DuckDBLoggerWriteLogEntryFun -> IO ()
forall a. StablePtr a -> IO ()
freeStablePtr StablePtr DuckDBLoggerWriteLogEntryFun
callbackStable IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> DuckDBDeleteCallback -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBDeleteCallback
deleteCb
    bracket c_duckdb_create_log_storage destroyLogStorage \DuckDBLogStorage
storage ->
        (IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Text -> (CString -> IO ()) -> IO ()
forall a. Text -> (CString -> IO a) -> IO a
TextForeign.withCString Text
name \CString
cName ->
                DuckDBLogStorage -> CString -> IO ()
c_duckdb_log_storage_set_name DuckDBLogStorage
storage CString
cName
            DuckDBLogStorage -> DuckDBLoggerWriteLogEntryFun -> IO ()
c_duckdb_log_storage_set_write_log_entry DuckDBLogStorage
storage DuckDBLoggerWriteLogEntryFun
writeCb
            DuckDBLogStorage -> Ptr () -> DuckDBDeleteCallback -> IO ()
c_duckdb_log_storage_set_extra_data DuckDBLogStorage
storage (StablePtr DuckDBLoggerWriteLogEntryFun -> Ptr ()
forall a. StablePtr a -> Ptr ()
castStablePtrToPtr StablePtr DuckDBLoggerWriteLogEntryFun
callbackStable) DuckDBDeleteCallback
deleteCb
            Connection -> (DuckDBDatabase -> IO ()) -> IO ()
forall a. Connection -> (DuckDBDatabase -> IO a) -> IO a
withDatabaseHandle Connection
conn \DuckDBDatabase
db -> do
                rc <- DuckDBDatabase -> DuckDBLogStorage -> IO DuckDBState
c_duckdb_register_log_storage DuckDBDatabase
db DuckDBLogStorage
storage
                if rc == DuckDBSuccess
                    then pure ()
                    else throwRegistrationError "register log storage"

logStorageHandler ::
    (LogEntry -> IO ()) ->
    Ptr () ->
    Ptr DuckDBTimestamp ->
    CString ->
    CString ->
    CString ->
    IO ()
logStorageHandler :: (LogEntry -> IO ())
-> Ptr ()
-> Ptr DuckDBTimestamp
-> CString
-> CString
-> CString
-> IO ()
logStorageHandler LogEntry -> IO ()
callback Ptr ()
_ Ptr DuckDBTimestamp
timestampPtr CString
levelPtr CString
logTypePtr CString
messagePtr = do
    entry <- do
        logEntryTimestamp <- Ptr DuckDBTimestamp -> IO (Maybe UTCTime)
readTimestamp Ptr DuckDBTimestamp
timestampPtr
        logEntryLevel <- readCStringText levelPtr
        logEntryType <- readCStringText logTypePtr
        logEntryMessage <- readCStringText messagePtr
        pure LogEntry{logEntryTimestamp, logEntryLevel, logEntryType, logEntryMessage}
    callback entry

readTimestamp :: Ptr DuckDBTimestamp -> IO (Maybe UTCTime)
readTimestamp :: Ptr DuckDBTimestamp -> IO (Maybe UTCTime)
readTimestamp Ptr DuckDBTimestamp
ptr
    | Ptr DuckDBTimestamp
ptr Ptr DuckDBTimestamp -> Ptr DuckDBTimestamp -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr DuckDBTimestamp
forall a. Ptr a
nullPtr = Maybe UTCTime -> IO (Maybe UTCTime)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe UTCTime
forall a. Maybe a
Nothing
    | Bool
otherwise = do
        DuckDBTimestamp micros <- Ptr DuckDBTimestamp -> IO DuckDBTimestamp
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBTimestamp
ptr
        pure (Just (posixSecondsToUTCTime (fromRational (toInteger micros % 1000000))))

readCStringText :: CString -> IO Text
readCStringText :: CString -> IO Text
readCStringText CString
ptr
    | CString
ptr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr = Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
Text.empty
    | Bool
otherwise = String -> Text
Text.pack (String -> Text) -> IO String -> IO Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CString -> IO String
peekCString CString
ptr

releaseWriteLogCallback :: Ptr () -> IO ()
releaseWriteLogCallback :: Ptr () -> IO ()
releaseWriteLogCallback Ptr ()
rawPtr =
    if Ptr ()
rawPtr Ptr () -> Ptr () -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr ()
forall a. Ptr a
nullPtr
        then () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        else do
            let stablePtr :: StablePtr DuckDBLoggerWriteLogEntryFun
stablePtr = Ptr () -> StablePtr DuckDBLoggerWriteLogEntryFun
forall a. Ptr () -> StablePtr a
castPtrToStablePtr Ptr ()
rawPtr :: StablePtr DuckDBLoggerWriteLogEntryFun
            callback <- StablePtr DuckDBLoggerWriteLogEntryFun
-> IO DuckDBLoggerWriteLogEntryFun
forall a. StablePtr a -> IO a
deRefStablePtr StablePtr DuckDBLoggerWriteLogEntryFun
stablePtr
            freeHaskellFunPtr callback
            freeStablePtr stablePtr

destroyLogStorage :: DuckDBLogStorage -> IO ()
destroyLogStorage :: DuckDBLogStorage -> IO ()
destroyLogStorage DuckDBLogStorage
storage =
    (Ptr DuckDBLogStorage -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBLogStorage
ptr -> Ptr DuckDBLogStorage -> DuckDBLogStorage -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBLogStorage
ptr DuckDBLogStorage
storage IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr DuckDBLogStorage -> IO ()
c_duckdb_destroy_log_storage Ptr DuckDBLogStorage
ptr

foreign import ccall "wrapper"
    mkWriteLogEntryCallback ::
        (Ptr () -> Ptr DuckDBTimestamp -> CString -> CString -> CString -> IO ()) ->
        IO DuckDBLoggerWriteLogEntryFun