{-# LANGUAGE BlockArguments #-}

{- |
Module      : Database.DuckDB.Simple.FileSystem
Description : High-level wrappers around DuckDB's file-system API.
-}
module Database.DuckDB.Simple.FileSystem (
    withFileHandle,
    readFileHandleChunk,
    writeFileHandleBytes,
    fileHandleTell,
    fileHandleSize,
    fileHandleSeek,
    fileHandleSync,
) where

import Control.Exception (bracket, throwIO)
import qualified Data.ByteString as BS
import Data.Int (Int64)
import Data.Text (Text)
import qualified Data.Text as Text
import Database.DuckDB.FFI
import Database.DuckDB.Simple.Internal (Connection, SQLError (..), withClientContext)
import Foreign.C.String (peekCString, withCString)
import Foreign.Marshal.Alloc (alloca, free, mallocBytes)
import Foreign.Ptr (castPtr, nullPtr)
import Foreign.Storable (peek, poke)

-- | Open a file through DuckDB's file-system layer for the duration of an action.
withFileHandle :: Connection -> FilePath -> [DuckDBFileFlag] -> (DuckDBFileHandle -> IO a) -> IO a
withFileHandle :: forall a.
Connection
-> String -> [DuckDBFileFlag] -> (DuckDBFileHandle -> IO a) -> IO a
withFileHandle Connection
conn String
path [DuckDBFileFlag]
flags DuckDBFileHandle -> IO a
action =
    Connection -> (DuckDBFileSystem -> IO a) -> IO a
forall a. Connection -> (DuckDBFileSystem -> IO a) -> IO a
withFileSystem Connection
conn \DuckDBFileSystem
fs ->
        IO DuckDBFileOpenOptions
-> (DuckDBFileOpenOptions -> IO ())
-> (DuckDBFileOpenOptions -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
            IO DuckDBFileOpenOptions
c_duckdb_create_file_open_options
            DuckDBFileOpenOptions -> IO ()
destroyFileOpenOptions
            \DuckDBFileOpenOptions
opts -> do
                (DuckDBFileFlag -> IO ()) -> [DuckDBFileFlag] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\DuckDBFileFlag
flag -> String -> IO DuckDBState -> IO ()
expectState String
"set file-open flag" (DuckDBFileOpenOptions -> DuckDBFileFlag -> CBool -> IO DuckDBState
c_duckdb_file_open_options_set_flag DuckDBFileOpenOptions
opts DuckDBFileFlag
flag CBool
1)) [DuckDBFileFlag]
flags
                String -> (CString -> IO a) -> IO a
forall a. String -> (CString -> IO a) -> IO a
withCString String
path \CString
cPath ->
                    (Ptr DuckDBFileHandle -> IO a) -> IO a
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBFileHandle
filePtr -> do
                        rc <- DuckDBFileSystem
-> CString
-> DuckDBFileOpenOptions
-> Ptr DuckDBFileHandle
-> IO DuckDBState
c_duckdb_file_system_open DuckDBFileSystem
fs CString
cPath DuckDBFileOpenOptions
opts Ptr DuckDBFileHandle
filePtr
                        if rc /= DuckDBSuccess
                            then throwFileSystemError fs path
                            else do
                                handle <- peek filePtr
                                bracket (pure handle) destroyFileHandle action

-- | Read up to the requested number of bytes from a file handle.
readFileHandleChunk :: DuckDBFileHandle -> Int64 -> IO BS.ByteString
readFileHandleChunk :: DuckDBFileHandle -> Int64 -> IO ByteString
readFileHandleChunk DuckDBFileHandle
handle Int64
requested
    | Int64
requested Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int64
0 = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
    | Bool
otherwise = do
        raw <- Int -> IO (Ptr ())
forall a. Int -> IO (Ptr a)
mallocBytes (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
requested)
        bytesRead <- c_duckdb_file_handle_read handle raw requested
        if bytesRead < 0
            then free raw >> throwFileHandleError handle (Text.pack "read failed")
            else do
                bs <- BS.packCStringLen (castPtr raw, fromIntegral bytesRead)
                free raw
                pure bs

-- | Write an entire bytestring to a file handle.
writeFileHandleBytes :: DuckDBFileHandle -> BS.ByteString -> IO Int64
writeFileHandleBytes :: DuckDBFileHandle -> ByteString -> IO Int64
writeFileHandleBytes DuckDBFileHandle
handle ByteString
bytes =
    ByteString -> (CStringLen -> IO Int64) -> IO Int64
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.useAsCStringLen ByteString
bytes \(CString
ptr, Int
len) -> do
        written <- DuckDBFileHandle -> Ptr () -> Int64 -> IO Int64
c_duckdb_file_handle_write DuckDBFileHandle
handle (CString -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr CString
ptr) (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
        if written < 0
            then throwFileHandleError handle (Text.pack "write failed")
            else pure written

-- | Return the current file position.
fileHandleTell :: DuckDBFileHandle -> IO Int64
fileHandleTell :: DuckDBFileHandle -> IO Int64
fileHandleTell DuckDBFileHandle
handle = do
    pos <- DuckDBFileHandle -> IO Int64
c_duckdb_file_handle_tell DuckDBFileHandle
handle
    if pos < 0 then throwFileHandleError handle (Text.pack "tell failed") else pure pos

-- | Return the current file size in bytes.
fileHandleSize :: DuckDBFileHandle -> IO Int64
fileHandleSize :: DuckDBFileHandle -> IO Int64
fileHandleSize DuckDBFileHandle
handle = do
    size <- DuckDBFileHandle -> IO Int64
c_duckdb_file_handle_size DuckDBFileHandle
handle
    if size < 0 then throwFileHandleError handle (Text.pack "size failed") else pure size

-- | Seek to an absolute byte offset.
fileHandleSeek :: DuckDBFileHandle -> Int64 -> IO ()
fileHandleSeek :: DuckDBFileHandle -> Int64 -> IO ()
fileHandleSeek DuckDBFileHandle
handle Int64
pos = do
    rc <- DuckDBFileHandle -> Int64 -> IO DuckDBState
c_duckdb_file_handle_seek DuckDBFileHandle
handle Int64
pos
    if rc == DuckDBSuccess
        then pure ()
        else throwFileHandleError handle (Text.pack "seek failed")

-- | Flush file-handle writes to stable storage.
fileHandleSync :: DuckDBFileHandle -> IO ()
fileHandleSync :: DuckDBFileHandle -> IO ()
fileHandleSync DuckDBFileHandle
handle = do
    rc <- DuckDBFileHandle -> IO DuckDBState
c_duckdb_file_handle_sync DuckDBFileHandle
handle
    if rc == DuckDBSuccess
        then pure ()
        else throwFileHandleError handle (Text.pack "sync failed")

withFileSystem :: Connection -> (DuckDBFileSystem -> IO a) -> IO a
withFileSystem :: forall a. Connection -> (DuckDBFileSystem -> IO a) -> IO a
withFileSystem Connection
conn DuckDBFileSystem -> IO a
action =
    Connection -> (DuckDBClientContext -> IO a) -> IO a
forall a. Connection -> (DuckDBClientContext -> IO a) -> IO a
withClientContext Connection
conn \DuckDBClientContext
ctx ->
        IO DuckDBFileSystem
-> (DuckDBFileSystem -> IO ())
-> (DuckDBFileSystem -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
            (DuckDBClientContext -> IO DuckDBFileSystem
c_duckdb_client_context_get_file_system DuckDBClientContext
ctx)
            DuckDBFileSystem -> IO ()
destroyFileSystem
            DuckDBFileSystem -> IO a
action

destroyFileSystem :: DuckDBFileSystem -> IO ()
destroyFileSystem :: DuckDBFileSystem -> IO ()
destroyFileSystem DuckDBFileSystem
fs =
    (Ptr DuckDBFileSystem -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBFileSystem
ptr -> Ptr DuckDBFileSystem -> DuckDBFileSystem -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBFileSystem
ptr DuckDBFileSystem
fs 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 DuckDBFileSystem -> IO ()
c_duckdb_destroy_file_system Ptr DuckDBFileSystem
ptr

destroyFileOpenOptions :: DuckDBFileOpenOptions -> IO ()
destroyFileOpenOptions :: DuckDBFileOpenOptions -> IO ()
destroyFileOpenOptions DuckDBFileOpenOptions
opts =
    (Ptr DuckDBFileOpenOptions -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBFileOpenOptions
ptr -> Ptr DuckDBFileOpenOptions -> DuckDBFileOpenOptions -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBFileOpenOptions
ptr DuckDBFileOpenOptions
opts 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 DuckDBFileOpenOptions -> IO ()
c_duckdb_destroy_file_open_options Ptr DuckDBFileOpenOptions
ptr

destroyFileHandle :: DuckDBFileHandle -> IO ()
destroyFileHandle :: DuckDBFileHandle -> IO ()
destroyFileHandle DuckDBFileHandle
handle =
    (Ptr DuckDBFileHandle -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBFileHandle
ptr -> Ptr DuckDBFileHandle -> DuckDBFileHandle -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBFileHandle
ptr DuckDBFileHandle
handle 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 DuckDBFileHandle -> IO ()
c_duckdb_destroy_file_handle Ptr DuckDBFileHandle
ptr

throwFileSystemError :: DuckDBFileSystem -> FilePath -> IO a
throwFileSystemError :: forall a. DuckDBFileSystem -> String -> IO a
throwFileSystemError DuckDBFileSystem
fs String
path = do
    err <- DuckDBFileSystem -> IO DuckDBErrorData
c_duckdb_file_system_error_data DuckDBFileSystem
fs
    throwErrorData err (Text.concat [Text.pack "duckdb-simple: failed to open file ", Text.pack path])

throwFileHandleError :: DuckDBFileHandle -> Text -> IO a
throwFileHandleError :: forall a. DuckDBFileHandle -> Text -> IO a
throwFileHandleError DuckDBFileHandle
handle Text
fallback = do
    err <- DuckDBFileHandle -> IO DuckDBErrorData
c_duckdb_file_handle_error_data DuckDBFileHandle
handle
    throwErrorData err fallback

throwErrorData :: DuckDBErrorData -> Text -> IO a
throwErrorData :: forall a. DuckDBErrorData -> Text -> IO a
throwErrorData DuckDBErrorData
err Text
fallback =
    IO DuckDBErrorData
-> (DuckDBErrorData -> IO ()) -> (DuckDBErrorData -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (DuckDBErrorData -> IO DuckDBErrorData
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DuckDBErrorData
err) DuckDBErrorData -> IO ()
destroyErrorData \DuckDBErrorData
errData -> do
        msgPtr <- DuckDBErrorData -> IO CString
c_duckdb_error_data_message DuckDBErrorData
errData
        errType <- c_duckdb_error_data_error_type errData
        message <-
            if msgPtr == nullPtr
                then pure fallback
                else Text.pack <$> peekCString msgPtr
        throwIO
            SQLError
                { sqlErrorMessage = message
                , sqlErrorType = Just errType
                , sqlErrorQuery = Nothing
                }

destroyErrorData :: DuckDBErrorData -> IO ()
destroyErrorData :: DuckDBErrorData -> IO ()
destroyErrorData DuckDBErrorData
err =
    (Ptr DuckDBErrorData -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBErrorData
ptr -> Ptr DuckDBErrorData -> DuckDBErrorData -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBErrorData
ptr DuckDBErrorData
err 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 DuckDBErrorData -> IO ()
c_duckdb_destroy_error_data Ptr DuckDBErrorData
ptr

expectState :: String -> IO DuckDBState -> IO ()
expectState :: String -> IO DuckDBState -> IO ()
expectState String
label IO DuckDBState
action = do
    rc <- IO DuckDBState
action
    if rc == DuckDBSuccess
        then pure ()
        else
            throwIO $
                SQLError
                    { sqlErrorMessage = Text.pack ("duckdb-simple: " <> label <> " failed")
                    , sqlErrorType = Nothing
                    , sqlErrorQuery = Nothing
                    }