{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

{- |
Module      : Database.DuckDB.Simple.Function
Description : Register scalar Haskell functions with DuckDB connections.

This module mirrors the high-level API provided by @sqlite-simple@ for user
defined functions, adapted to DuckDB's chunked execution model.  It allows
pure and 'IO'-based Haskell functions to be exposed to SQL while reusing the
existing 'FromField' and 'ToField'-style typeclass machinery for argument
decoding and result marshalling.
-}
module Database.DuckDB.Simple.Function (
    Function,
    createFunction,
    deleteFunction,
) where

import Control.Exception (
    SomeException,
    bracket,
    displayException,
    onException,
    throwIO,
    try,
 )
import Control.Monad (forM, forM_, when)
import qualified Data.ByteString as BS
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Proxy (Proxy (..))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Foreign as TextForeign
import Data.Word (Word16, Word32, Word64, Word8)
import Database.DuckDB.FFI
import Database.DuckDB.Simple.FromField (
    Field (..),
    FieldValue (..),
    FromField (..),
 )
import Database.DuckDB.Simple.Internal (
    Connection,
    Query (..),
    SQLError (..),
    withConnectionHandle,
    withQueryCString,
 )
import Foreign.C.String (peekCString)
import Foreign.C.Types (CBool (..))
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr, castPtr, freeHaskellFunPtr, nullPtr, plusPtr)
import Foreign.StablePtr (StablePtr, castPtrToStablePtr, castStablePtrToPtr, deRefStablePtr, freeStablePtr, newStablePtr)
import Foreign.Storable (peekElemOff, poke, pokeElemOff)
import Database.DuckDB.Simple.Ok (Ok(..))

-- | Tag DuckDB logical types we support for scalar return values.
data ScalarType
    = ScalarTypeBoolean
    | ScalarTypeBigInt
    | ScalarTypeDouble
    | ScalarTypeVarchar

-- | Runtime representation of values returned to DuckDB.
data ScalarValue
    = ScalarNull
    | ScalarBoolean !Bool
    | ScalarInteger !Int64
    | ScalarDouble !Double
    | ScalarText !Text

-- | Class of scalar results that can be produced by user-defined functions.
class FunctionResult a where
    scalarReturnType :: Proxy a -> ScalarType
    toScalarValue :: a -> IO ScalarValue

instance FunctionResult Int where
    scalarReturnType :: Proxy Int -> ScalarType
scalarReturnType Proxy Int
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Int -> IO ScalarValue
toScalarValue Int
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
value))

instance FunctionResult Int16 where
    scalarReturnType :: Proxy Int16 -> ScalarType
scalarReturnType Proxy Int16
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Int16 -> IO ScalarValue
toScalarValue Int16
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Int16 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
value))

instance FunctionResult Int32 where
    scalarReturnType :: Proxy Int32 -> ScalarType
scalarReturnType Proxy Int32
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Int32 -> IO ScalarValue
toScalarValue Int32
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Int32 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
value))

instance FunctionResult Int64 where
    scalarReturnType :: Proxy Int64 -> ScalarType
scalarReturnType Proxy Int64
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Int64 -> IO ScalarValue
toScalarValue Int64
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger Int64
value)

instance FunctionResult Word where
    scalarReturnType :: Proxy Word -> ScalarType
scalarReturnType Proxy Word
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Word -> IO ScalarValue
toScalarValue Word
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Word -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
value))

instance FunctionResult Word16 where
    scalarReturnType :: Proxy Word16 -> ScalarType
scalarReturnType Proxy Word16
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Word16 -> IO ScalarValue
toScalarValue Word16
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Word16 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
value))

instance FunctionResult Word32 where
    scalarReturnType :: Proxy Word32 -> ScalarType
scalarReturnType Proxy Word32
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Word32 -> IO ScalarValue
toScalarValue Word32
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Word32 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
value))

instance FunctionResult Word64 where
    scalarReturnType :: Proxy Word64 -> ScalarType
scalarReturnType Proxy Word64
_ = ScalarType
ScalarTypeBigInt
    toScalarValue :: Word64 -> IO ScalarValue
toScalarValue Word64
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> ScalarValue
ScalarInteger (Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
value))

instance FunctionResult Double where
    scalarReturnType :: Proxy Double -> ScalarType
scalarReturnType Proxy Double
_ = ScalarType
ScalarTypeDouble
    toScalarValue :: Double -> IO ScalarValue
toScalarValue Double
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> ScalarValue
ScalarDouble Double
value)

instance FunctionResult Float where
    scalarReturnType :: Proxy Float -> ScalarType
scalarReturnType Proxy Float
_ = ScalarType
ScalarTypeDouble
    toScalarValue :: Float -> IO ScalarValue
toScalarValue Float
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> ScalarValue
ScalarDouble (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
value))

instance FunctionResult Bool where
    scalarReturnType :: Proxy Bool -> ScalarType
scalarReturnType Proxy Bool
_ = ScalarType
ScalarTypeBoolean
    toScalarValue :: Bool -> IO ScalarValue
toScalarValue Bool
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> ScalarValue
ScalarBoolean Bool
value)

instance FunctionResult Text where
    scalarReturnType :: Proxy Text -> ScalarType
scalarReturnType Proxy Text
_ = ScalarType
ScalarTypeVarchar
    toScalarValue :: Text -> IO ScalarValue
toScalarValue Text
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> ScalarValue
ScalarText Text
value)

instance FunctionResult String where
    scalarReturnType :: Proxy String -> ScalarType
scalarReturnType Proxy String
_ = ScalarType
ScalarTypeVarchar
    toScalarValue :: String -> IO ScalarValue
toScalarValue String
value = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> ScalarValue
ScalarText (String -> Text
Text.pack String
value))

instance (FunctionResult a) => FunctionResult (Maybe a) where
    scalarReturnType :: Proxy (Maybe a) -> ScalarType
scalarReturnType Proxy (Maybe a)
_ = Proxy a -> ScalarType
forall a. FunctionResult a => Proxy a -> ScalarType
scalarReturnType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    toScalarValue :: Maybe a -> IO ScalarValue
toScalarValue Maybe a
Nothing = ScalarValue -> IO ScalarValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarValue
ScalarNull
    toScalarValue (Just a
value) = a -> IO ScalarValue
forall a. FunctionResult a => a -> IO ScalarValue
toScalarValue a
value

-- | Argument types supported by the scalar function machinery.
class FunctionArg a where
    argumentType :: Proxy a -> DuckDBType

instance FunctionArg Int where
    argumentType :: Proxy Int -> DuckDBType
argumentType Proxy Int
_ = DuckDBType
DuckDBTypeBigInt

instance FunctionArg Int16 where
    argumentType :: Proxy Int16 -> DuckDBType
argumentType Proxy Int16
_ = DuckDBType
DuckDBTypeSmallInt

instance FunctionArg Int32 where
    argumentType :: Proxy Int32 -> DuckDBType
argumentType Proxy Int32
_ = DuckDBType
DuckDBTypeInteger

instance FunctionArg Int64 where
    argumentType :: Proxy Int64 -> DuckDBType
argumentType Proxy Int64
_ = DuckDBType
DuckDBTypeBigInt

instance FunctionArg Word where
    argumentType :: Proxy Word -> DuckDBType
argumentType Proxy Word
_ = DuckDBType
DuckDBTypeUBigInt

instance FunctionArg Word16 where
    argumentType :: Proxy Word16 -> DuckDBType
argumentType Proxy Word16
_ = DuckDBType
DuckDBTypeUSmallInt

instance FunctionArg Word32 where
    argumentType :: Proxy Word32 -> DuckDBType
argumentType Proxy Word32
_ = DuckDBType
DuckDBTypeUInteger

instance FunctionArg Word64 where
    argumentType :: Proxy Word64 -> DuckDBType
argumentType Proxy Word64
_ = DuckDBType
DuckDBTypeUBigInt

instance FunctionArg Double where
    argumentType :: Proxy Double -> DuckDBType
argumentType Proxy Double
_ = DuckDBType
DuckDBTypeDouble

instance FunctionArg Float where
    argumentType :: Proxy Float -> DuckDBType
argumentType Proxy Float
_ = DuckDBType
DuckDBTypeFloat

instance FunctionArg Bool where
    argumentType :: Proxy Bool -> DuckDBType
argumentType Proxy Bool
_ = DuckDBType
DuckDBTypeBoolean

instance FunctionArg Text where
    argumentType :: Proxy Text -> DuckDBType
argumentType Proxy Text
_ = DuckDBType
DuckDBTypeVarchar

instance FunctionArg String where
    argumentType :: Proxy String -> DuckDBType
argumentType Proxy String
_ = DuckDBType
DuckDBTypeVarchar

instance (FunctionArg a) => FunctionArg (Maybe a) where
    argumentType :: Proxy (Maybe a) -> DuckDBType
argumentType Proxy (Maybe a)
_ = Proxy a -> DuckDBType
forall a. FunctionArg a => Proxy a -> DuckDBType
argumentType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)

-- | Typeclass describing Haskell functions that can be exposed to DuckDB.
class Function a where
    argumentTypes :: Proxy a -> [DuckDBType]
    returnType :: Proxy a -> ScalarType
    isVolatile :: Proxy a -> Bool
    applyFunction :: [Field] -> a -> IO ScalarValue

instance {-# OVERLAPPABLE #-} (FunctionResult a) => Function a where
    argumentTypes :: Proxy a -> [DuckDBType]
argumentTypes Proxy a
_ = []
    returnType :: Proxy a -> ScalarType
returnType Proxy a
_ = Proxy a -> ScalarType
forall a. FunctionResult a => Proxy a -> ScalarType
scalarReturnType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    isVolatile :: Proxy a -> Bool
isVolatile Proxy a
_ = Bool
False
    applyFunction :: [Field] -> a -> IO ScalarValue
applyFunction [] a
value = a -> IO ScalarValue
forall a. FunctionResult a => a -> IO ScalarValue
toScalarValue a
value
    applyFunction [Field]
_ a
_ = SQLError -> IO ScalarValue
forall e a. Exception e => e -> IO a
throwIO (Text -> SQLError
functionInvocationError (String -> Text
Text.pack String
"unexpected arguments supplied"))

instance {-# OVERLAPPING #-} (FunctionResult a) => Function (IO a) where
    argumentTypes :: Proxy (IO a) -> [DuckDBType]
argumentTypes Proxy (IO a)
_ = []
    returnType :: Proxy (IO a) -> ScalarType
returnType Proxy (IO a)
_ = Proxy a -> ScalarType
forall a. FunctionResult a => Proxy a -> ScalarType
scalarReturnType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    isVolatile :: Proxy (IO a) -> Bool
isVolatile Proxy (IO a)
_ = Bool
True
    applyFunction :: [Field] -> IO a -> IO ScalarValue
applyFunction [] IO a
action = IO a
action IO a -> (a -> IO ScalarValue) -> IO ScalarValue
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> IO ScalarValue
forall a. FunctionResult a => a -> IO ScalarValue
toScalarValue
    applyFunction [Field]
_ IO a
_ = SQLError -> IO ScalarValue
forall e a. Exception e => e -> IO a
throwIO (Text -> SQLError
functionInvocationError (String -> Text
Text.pack String
"unexpected arguments supplied"))

instance {-# OVERLAPPABLE #-} (FromField a, FunctionArg a, Function r) => Function (a -> r) where
    argumentTypes :: Proxy (a -> r) -> [DuckDBType]
argumentTypes Proxy (a -> r)
_ = Proxy a -> DuckDBType
forall a. FunctionArg a => Proxy a -> DuckDBType
argumentType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a) DuckDBType -> [DuckDBType] -> [DuckDBType]
forall a. a -> [a] -> [a]
: Proxy r -> [DuckDBType]
forall a. Function a => Proxy a -> [DuckDBType]
argumentTypes (Proxy r
forall {k} (t :: k). Proxy t
Proxy :: Proxy r)
    returnType :: Proxy (a -> r) -> ScalarType
returnType Proxy (a -> r)
_ = Proxy r -> ScalarType
forall a. Function a => Proxy a -> ScalarType
returnType (Proxy r
forall {k} (t :: k). Proxy t
Proxy :: Proxy r)
    isVolatile :: Proxy (a -> r) -> Bool
isVolatile Proxy (a -> r)
_ = Proxy r -> Bool
forall a. Function a => Proxy a -> Bool
isVolatile (Proxy r
forall {k} (t :: k). Proxy t
Proxy :: Proxy r)
    applyFunction :: [Field] -> (a -> r) -> IO ScalarValue
applyFunction [] a -> r
_ =
        SQLError -> IO ScalarValue
forall e a. Exception e => e -> IO a
throwIO (Text -> SQLError
functionInvocationError (String -> Text
Text.pack String
"insufficient arguments supplied"))
    applyFunction (Field
field : [Field]
rest) a -> r
fn =
        case FieldParser a
forall a. FromField a => FieldParser a
fromField Field
field of
            Errors [SomeException]
err -> SQLError -> IO ScalarValue
forall e a. Exception e => e -> IO a
throwIO (Int -> [SomeException] -> SQLError
argumentConversionError (Field -> Int
fieldIndex Field
field) [SomeException]
err)
            Ok a
value -> [Field] -> r -> IO ScalarValue
forall a. Function a => [Field] -> a -> IO ScalarValue
applyFunction [Field]
rest (a -> r
fn a
value)

-- | Register a Haskell function under the supplied name.
createFunction :: forall f. (Function f) => Connection -> Text -> f -> IO ()
createFunction :: forall f. Function f => Connection -> Text -> f -> IO ()
createFunction Connection
conn Text
name f
fn = do
    DuckDBScalarFunctionFun
funPtr <- (DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ())
-> IO DuckDBScalarFunctionFun
mkScalarFun (f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
forall f.
Function f =>
f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandler f
fn)
    StablePtr DuckDBScalarFunctionFun
funPtrStable <- DuckDBScalarFunctionFun -> IO (StablePtr DuckDBScalarFunctionFun)
forall a. a -> IO (StablePtr a)
newStablePtr DuckDBScalarFunctionFun
funPtr
    DuckDBDeleteCallback
destroyCb <- (Ptr () -> IO ()) -> IO DuckDBDeleteCallback
mkDeleteCallback Ptr () -> IO ()
releaseFunctionPtr
    let release :: IO ()
release = DuckDBScalarFunctionFun
-> StablePtr DuckDBScalarFunctionFun
-> DuckDBDeleteCallback
-> IO ()
destroyFunctionResources DuckDBScalarFunctionFun
funPtr StablePtr DuckDBScalarFunctionFun
funPtrStable DuckDBDeleteCallback
destroyCb
    IO DuckDBScalarFunction
-> (DuckDBScalarFunction -> IO ())
-> (DuckDBScalarFunction -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO DuckDBScalarFunction
c_duckdb_create_scalar_function DuckDBScalarFunction -> IO ()
cleanupScalarFunction \DuckDBScalarFunction
scalarFun ->
        (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 -> IO ()) -> IO ()) -> (CString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CString
cName ->
                DuckDBScalarFunction -> CString -> IO ()
c_duckdb_scalar_function_set_name DuckDBScalarFunction
scalarFun CString
cName
            [DuckDBType] -> (DuckDBType -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Proxy f -> [DuckDBType]
forall a. Function a => Proxy a -> [DuckDBType]
argumentTypes (Proxy f
forall {k} (t :: k). Proxy t
Proxy :: Proxy f)) \DuckDBType
dtype ->
                DuckDBType -> (DuckDBLogicalType -> IO ()) -> IO ()
forall a. DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a
withLogicalType DuckDBType
dtype ((DuckDBLogicalType -> IO ()) -> IO ())
-> (DuckDBLogicalType -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \DuckDBLogicalType
logical ->
                    DuckDBScalarFunction -> DuckDBLogicalType -> IO ()
c_duckdb_scalar_function_add_parameter DuckDBScalarFunction
scalarFun DuckDBLogicalType
logical
            DuckDBType -> (DuckDBLogicalType -> IO ()) -> IO ()
forall a. DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a
withLogicalType (ScalarType -> DuckDBType
duckTypeForScalar (Proxy f -> ScalarType
forall a. Function a => Proxy a -> ScalarType
returnType (Proxy f
forall {k} (t :: k). Proxy t
Proxy :: Proxy f))) ((DuckDBLogicalType -> IO ()) -> IO ())
-> (DuckDBLogicalType -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \DuckDBLogicalType
logical ->
                DuckDBScalarFunction -> DuckDBLogicalType -> IO ()
c_duckdb_scalar_function_set_return_type DuckDBScalarFunction
scalarFun DuckDBLogicalType
logical
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Proxy f -> Bool
forall a. Function a => Proxy a -> Bool
isVolatile (Proxy f
forall {k} (t :: k). Proxy t
Proxy :: Proxy f)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                DuckDBScalarFunction -> IO ()
c_duckdb_scalar_function_set_volatile DuckDBScalarFunction
scalarFun
            DuckDBScalarFunction -> DuckDBScalarFunctionFun -> IO ()
c_duckdb_scalar_function_set_function DuckDBScalarFunction
scalarFun DuckDBScalarFunctionFun
funPtr
            DuckDBScalarFunction -> Ptr () -> DuckDBDeleteCallback -> IO ()
c_duckdb_scalar_function_set_extra_info DuckDBScalarFunction
scalarFun (StablePtr DuckDBScalarFunctionFun -> Ptr ()
forall a. StablePtr a -> Ptr ()
castStablePtrToPtr StablePtr DuckDBScalarFunctionFun
funPtrStable) DuckDBDeleteCallback
destroyCb
            Connection -> (DuckDBConnection -> IO ()) -> IO ()
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr -> do
                DuckDBState
rc <- DuckDBConnection -> DuckDBScalarFunction -> IO DuckDBState
c_duckdb_register_scalar_function DuckDBConnection
connPtr DuckDBScalarFunction
scalarFun
                if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                    then () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                    else SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (Text -> SQLError
functionInvocationError (String -> Text
Text.pack String
"duckdb-simple: registering function failed"))

-- | Drop a previously registered scalar function by issuing a DROP FUNCTION statement.
deleteFunction :: Connection -> Text -> IO ()
deleteFunction :: Connection -> Text -> IO ()
deleteFunction Connection
conn Text
name =
    do
        Either SQLError ()
outcome <-
            IO () -> IO (Either SQLError ())
forall e a. Exception e => IO a -> IO (Either e a)
try (IO () -> IO (Either SQLError ()))
-> IO () -> IO (Either SQLError ())
forall a b. (a -> b) -> a -> b
$
                Connection -> (DuckDBConnection -> IO ()) -> IO ()
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr -> do
                    let dropQuery :: Query
dropQuery =
                            Text -> Query
Query (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$
                                [Text] -> Text
Text.concat
                                    [ String -> Text
Text.pack String
"DROP FUNCTION IF EXISTS "
                                    , Text -> Text
qualifyIdentifier Text
name
                                    ]
                    Query -> (CString -> IO ()) -> IO ()
forall a. Query -> (CString -> IO a) -> IO a
withQueryCString Query
dropQuery \CString
sql ->
                        (Ptr DuckDBResult -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBResult
resPtr -> do
                            DuckDBState
rc <- DuckDBConnection -> CString -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_query DuckDBConnection
connPtr CString
sql Ptr DuckDBResult
resPtr
                            if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                                then Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                                else do
                                    Text
errMsg <- Ptr DuckDBResult -> IO Text
fetchResultError Ptr DuckDBResult
resPtr
                                    Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                                    SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO
                                        SQLError
                                            { sqlErrorMessage :: Text
sqlErrorMessage = Text
errMsg
                                            , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
forall a. Maybe a
Nothing
                                            , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Query -> Maybe Query
forall a. a -> Maybe a
Just Query
dropQuery
                                            }
        case Either SQLError ()
outcome of
            Right () -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Left SQLError
err
                -- DuckDB does not allow dropping scalar functions registered via the C API,
                -- so we ignore that specific error here.
                -- TODO: Update this when DuckDB adds support for dropping such functions.
                | Text -> Text -> Bool
Text.isInfixOf (String -> Text
Text.pack String
"Cannot drop internal catalog entry") (SQLError -> Text
sqlErrorMessage SQLError
err) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                | Bool
otherwise -> SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO SQLError
err

cleanupScalarFunction :: DuckDBScalarFunction -> IO ()
cleanupScalarFunction :: DuckDBScalarFunction -> IO ()
cleanupScalarFunction DuckDBScalarFunction
scalarFun =
    (Ptr DuckDBScalarFunction -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBScalarFunction
ptr -> do
        Ptr DuckDBScalarFunction -> DuckDBScalarFunction -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBScalarFunction
ptr DuckDBScalarFunction
scalarFun
        Ptr DuckDBScalarFunction -> IO ()
c_duckdb_destroy_scalar_function Ptr DuckDBScalarFunction
ptr

destroyFunctionResources :: DuckDBScalarFunctionFun -> StablePtr DuckDBScalarFunctionFun -> DuckDBDeleteCallback -> IO ()
destroyFunctionResources :: DuckDBScalarFunctionFun
-> StablePtr DuckDBScalarFunctionFun
-> DuckDBDeleteCallback
-> IO ()
destroyFunctionResources DuckDBScalarFunctionFun
funPtr StablePtr DuckDBScalarFunctionFun
funPtrStable DuckDBDeleteCallback
destroyCb = do
    DuckDBScalarFunctionFun -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBScalarFunctionFun
funPtr
    StablePtr DuckDBScalarFunctionFun -> IO ()
forall a. StablePtr a -> IO ()
freeStablePtr StablePtr DuckDBScalarFunctionFun
funPtrStable
    DuckDBDeleteCallback -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBDeleteCallback
destroyCb

releaseFunctionPtr :: Ptr () -> IO ()
releaseFunctionPtr :: Ptr () -> IO ()
releaseFunctionPtr Ptr ()
rawPtr =
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ptr ()
rawPtr Ptr () -> Ptr () -> Bool
forall a. Eq a => a -> a -> Bool
/= Ptr ()
forall a. Ptr a
nullPtr) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        let stablePtr :: StablePtr DuckDBScalarFunctionFun
stablePtr = Ptr () -> StablePtr DuckDBScalarFunctionFun
forall a. Ptr () -> StablePtr a
castPtrToStablePtr Ptr ()
rawPtr :: StablePtr DuckDBScalarFunctionFun
        DuckDBScalarFunctionFun
funPtr <- StablePtr DuckDBScalarFunctionFun -> IO DuckDBScalarFunctionFun
forall a. StablePtr a -> IO a
deRefStablePtr StablePtr DuckDBScalarFunctionFun
stablePtr
        DuckDBScalarFunctionFun -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBScalarFunctionFun
funPtr
        StablePtr DuckDBScalarFunctionFun -> IO ()
forall a. StablePtr a -> IO ()
freeStablePtr StablePtr DuckDBScalarFunctionFun
stablePtr

withLogicalType :: DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a
withLogicalType :: forall a. DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a
withLogicalType DuckDBType
dtype =
    IO DuckDBLogicalType
-> (DuckDBLogicalType -> IO ())
-> (DuckDBLogicalType -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
        ( do
            DuckDBLogicalType
logical <- DuckDBType -> IO DuckDBLogicalType
c_duckdb_create_logical_type DuckDBType
dtype
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (DuckDBLogicalType
logical DuckDBLogicalType -> DuckDBLogicalType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBLogicalType
forall a. Ptr a
nullPtr) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO ()) -> SQLError -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Text -> SQLError
functionInvocationError (String -> Text
Text.pack String
"duckdb-simple: failed to allocate logical type")
            DuckDBLogicalType -> IO DuckDBLogicalType
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DuckDBLogicalType
logical
        )
        DuckDBLogicalType -> IO ()
destroyLogicalType

destroyLogicalType :: DuckDBLogicalType -> IO ()
destroyLogicalType :: DuckDBLogicalType -> IO ()
destroyLogicalType DuckDBLogicalType
logicalType =
    (Ptr DuckDBLogicalType -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBLogicalType
ptr -> do
        Ptr DuckDBLogicalType -> DuckDBLogicalType -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBLogicalType
ptr DuckDBLogicalType
logicalType
        Ptr DuckDBLogicalType -> IO ()
c_duckdb_destroy_logical_type Ptr DuckDBLogicalType
ptr

duckTypeForScalar :: ScalarType -> DuckDBType
duckTypeForScalar :: ScalarType -> DuckDBType
duckTypeForScalar = \case
    ScalarType
ScalarTypeBoolean -> DuckDBType
DuckDBTypeBoolean
    ScalarType
ScalarTypeBigInt -> DuckDBType
DuckDBTypeBigInt
    ScalarType
ScalarTypeDouble -> DuckDBType
DuckDBTypeDouble
    ScalarType
ScalarTypeVarchar -> DuckDBType
DuckDBTypeVarchar

scalarFunctionHandler :: forall f. (Function f) => f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandler :: forall f.
Function f =>
f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandler f
fn DuckDBFunctionInfo
info DuckDBDataChunk
chunk DuckDBVector
outVec = do
    Either SomeException ()
result <-
        IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try do
            Word64
rawColumnCount <- DuckDBDataChunk -> IO Word64
c_duckdb_data_chunk_get_column_count DuckDBDataChunk
chunk
            let columnCount :: Int
columnCount = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rawColumnCount :: Int
                expected :: Int
expected = [DuckDBType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Proxy f -> [DuckDBType]
forall a. Function a => Proxy a -> [DuckDBType]
argumentTypes (Proxy f
forall {k} (t :: k). Proxy t
Proxy :: Proxy f))
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
columnCount Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
expected) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO ()) -> SQLError -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Text -> SQLError
functionInvocationError (Text -> SQLError) -> Text -> SQLError
forall a b. (a -> b) -> a -> b
$
                        [Text] -> Text
Text.concat
                            [ String -> Text
Text.pack String
"duckdb-simple: function expected "
                            , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
expected)
                            , String -> Text
Text.pack String
" arguments but received "
                            , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
columnCount)
                            ]
            Word64
rawRowCount <- DuckDBDataChunk -> IO Word64
c_duckdb_data_chunk_get_size DuckDBDataChunk
chunk
            let rowCount :: Int
rowCount = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rawRowCount :: Int
            [ColumnReader]
readers <- (Int -> IO ColumnReader) -> [Int] -> IO [ColumnReader]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader DuckDBDataChunk
chunk) [Int
0 .. Int
expected Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            [[Field]]
rows <-
                [Int] -> (Int -> IO [Field]) -> IO [[Field]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int
rowCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] \Int
row ->
                    [ColumnReader] -> (ColumnReader -> IO Field) -> IO [Field]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [ColumnReader]
readers \ColumnReader
reader ->
                        ColumnReader
reader (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
row)
            [ScalarValue]
results <- ([Field] -> IO ScalarValue) -> [[Field]] -> IO [ScalarValue]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Field] -> f -> IO ScalarValue
forall a. Function a => [Field] -> a -> IO ScalarValue
`applyFunction` f
fn) [[Field]]
rows
            ScalarType -> [ScalarValue] -> DuckDBVector -> IO ()
writeResults (Proxy f -> ScalarType
forall a. Function a => Proxy a -> ScalarType
returnType (Proxy f
forall {k} (t :: k). Proxy t
Proxy :: Proxy f)) [ScalarValue]
results DuckDBVector
outVec
            DuckDBDataChunk -> Word64 -> IO ()
c_duckdb_data_chunk_set_size DuckDBDataChunk
chunk (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rowCount)
    case Either SomeException ()
result of
        Left (SomeException
err :: SomeException) -> do
            DuckDBDataChunk -> Word64 -> IO ()
c_duckdb_data_chunk_set_size DuckDBDataChunk
chunk Word64
0
            let message :: Text
message = String -> Text
Text.pack (SomeException -> String
forall e. Exception e => e -> String
displayException SomeException
err)
            Text -> (CString -> IO ()) -> IO ()
forall a. Text -> (CString -> IO a) -> IO a
TextForeign.withCString Text
message ((CString -> IO ()) -> IO ()) -> (CString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CString
cMsg ->
                DuckDBFunctionInfo -> CString -> IO ()
c_duckdb_scalar_function_set_error DuckDBFunctionInfo
info CString
cMsg
        Right () -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

type ColumnReader = DuckDBIdx -> IO Field

makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader DuckDBDataChunk
chunk Int
columnIndex = do
    DuckDBVector
vector <- DuckDBDataChunk -> Word64 -> IO DuckDBVector
c_duckdb_data_chunk_get_vector DuckDBDataChunk
chunk (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
columnIndex)
    DuckDBLogicalType
logical <- DuckDBVector -> IO DuckDBLogicalType
c_duckdb_vector_get_column_type DuckDBVector
vector
    DuckDBType
dtype <- DuckDBLogicalType -> IO DuckDBType
c_duckdb_get_type_id DuckDBLogicalType
logical
    DuckDBLogicalType -> IO ()
destroyLogicalType DuckDBLogicalType
logical
    Ptr ()
dataPtr <- DuckDBVector -> IO (Ptr ())
c_duckdb_vector_get_data DuckDBVector
vector
    Ptr Word64
validity <- DuckDBVector -> IO (Ptr Word64)
c_duckdb_vector_get_validity DuckDBVector
vector
    let name :: Text
name = String -> Text
Text.pack (String
"arg" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
columnIndex)
    ColumnReader -> IO ColumnReader
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure \Word64
rowIdx -> do
        Bool
valid <- Ptr Word64 -> Word64 -> IO Bool
isRowValid Ptr Word64
validity Word64
rowIdx
        FieldValue
value <-
            if Bool
valid
                then DuckDBType -> Ptr () -> Word64 -> IO FieldValue
fetchValue DuckDBType
dtype Ptr ()
dataPtr Word64
rowIdx
                else FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FieldValue
FieldNull
        Field -> IO Field
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            Field
                { fieldName :: Text
fieldName = Text
name
                , fieldIndex :: Int
fieldIndex = Int
columnIndex
                , fieldValue :: FieldValue
fieldValue = FieldValue
value
                }

isRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool
isRowValid :: Ptr Word64 -> Word64 -> IO Bool
isRowValid Ptr Word64
validity Word64
rowIdx
    | Ptr Word64
validity Ptr Word64 -> Ptr Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word64
forall a. Ptr a
nullPtr = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
    | Bool
otherwise = do
        CBool Word8
flag <- Ptr Word64 -> Word64 -> IO CBool
c_duckdb_validity_row_is_valid Ptr Word64
validity Word64
rowIdx
        Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word8
flag Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0)

fetchValue :: DuckDBType -> Ptr () -> DuckDBIdx -> IO FieldValue
fetchValue :: DuckDBType -> Ptr () -> Word64 -> IO FieldValue
fetchValue DuckDBType
dtype Ptr ()
dataPtr Word64
rowIdx
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeBoolean = do
        Word8
value <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> FieldValue
FieldBool (Word8
value Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeTinyInt = do
        Int8
value <- Ptr Int8 -> Int -> IO Int8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Int8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Int8) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int8 -> FieldValue
FieldInt8 (Int8 -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeSmallInt = do
        Int16
value <- Ptr Int16 -> Int -> IO Int16
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Int16
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Int16) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int16 -> FieldValue
FieldInt16 (Int16 -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int16
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeInteger = do
        Int32
value <- Ptr Int32 -> Int -> IO Int32
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Int32
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Int32) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int32 -> FieldValue
FieldInt32 (Int32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeBigInt Bool -> Bool -> Bool
|| DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeHugeInt = do
        Int64
value <- Ptr Int64 -> Int -> IO Int64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Int64
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Int64) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> FieldValue
FieldInt64 Int64
value)
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeUTinyInt = do
        Word8
value <- Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word8 -> FieldValue
FieldWord8 (Word8 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeUSmallInt = do
        Word16
value <- Ptr Word16 -> Int -> IO Word16
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Word16
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word16) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word16 -> FieldValue
FieldWord16 (Word16 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeUInteger = do
        Word32
value <- Ptr Word32 -> Int -> IO Word32
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Word32
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word32) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word32 -> FieldValue
FieldWord32 (Word32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeUBigInt Bool -> Bool -> Bool
|| DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeUHugeInt = do
        Word64
value <- Ptr Word64 -> Int -> IO Word64
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Word64
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word64) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word64 -> FieldValue
FieldWord64 (Word64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeFloat = do
        Float
value <- Ptr Float -> Int -> IO Float
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Float
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Float) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> FieldValue
FieldDouble (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
value))
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeDouble = do
        Double
value <- Ptr Double -> Int -> IO Double
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr Double
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Double) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> FieldValue
FieldDouble Double
value)
    | DuckDBType
dtype DuckDBType -> DuckDBType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBType
DuckDBTypeVarchar = do
        Text
textValue <- Ptr () -> Word64 -> IO Text
decodeDuckText Ptr ()
dataPtr Word64
rowIdx
        FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> FieldValue
FieldText Text
textValue)
    | Bool
otherwise =
        SQLError -> IO FieldValue
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO FieldValue) -> SQLError -> IO FieldValue
forall a b. (a -> b) -> a -> b
$
            Text -> SQLError
functionInvocationError (Text -> SQLError) -> Text -> SQLError
forall a b. (a -> b) -> a -> b
$
                [Text] -> Text
Text.concat
                    [ String -> Text
Text.pack String
"duckdb-simple: unsupported argument type "
                    , String -> Text
Text.pack (DuckDBType -> String
forall a. Show a => a -> String
show DuckDBType
dtype)
                    ]

decodeDuckText :: Ptr () -> DuckDBIdx -> IO Text
decodeDuckText :: Ptr () -> Word64 -> IO Text
decodeDuckText Ptr ()
dataPtr Word64
rowIdx = do
    let bytePtr :: Ptr Word8
bytePtr = Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8
        offset :: Int
offset = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
duckdbStringTSize
        stringPtr :: Ptr DuckDBStringT
stringPtr = Ptr Any -> Ptr DuckDBStringT
forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8
bytePtr Ptr Word8 -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
offset) :: Ptr DuckDBStringT
    Word32
len <- Ptr DuckDBStringT -> IO Word32
c_duckdb_string_t_length Ptr DuckDBStringT
stringPtr
    if Word32
len Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0
        then Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
Text.empty
        else do
            CString
cstr <- Ptr DuckDBStringT -> IO CString
c_duckdb_string_t_data Ptr DuckDBStringT
stringPtr
            ByteString
bytes <- CStringLen -> IO ByteString
BS.packCStringLen (CString
cstr, Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
            Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Text
TE.decodeUtf8 ByteString
bytes)

duckdbStringTSize :: Int
duckdbStringTSize :: Int
duckdbStringTSize = Int
16

writeResults :: ScalarType -> [ScalarValue] -> DuckDBVector -> IO ()
writeResults :: ScalarType -> [ScalarValue] -> DuckDBVector -> IO ()
writeResults ScalarType
resultType [ScalarValue]
values DuckDBVector
outVec = do
    let hasNulls :: Bool
hasNulls = (ScalarValue -> Bool) -> [ScalarValue] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ScalarValue -> Bool
isNullValue [ScalarValue]
values
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
hasNulls (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        DuckDBVector -> IO ()
c_duckdb_vector_ensure_validity_writable DuckDBVector
outVec
    Ptr ()
dataPtr <- DuckDBVector -> IO (Ptr ())
c_duckdb_vector_get_data DuckDBVector
outVec
    Ptr Word64
validityPtr <- DuckDBVector -> IO (Ptr Word64)
c_duckdb_vector_get_validity DuckDBVector
outVec
    [(Int, ScalarValue)] -> ((Int, ScalarValue) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [ScalarValue] -> [(Int, ScalarValue)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [ScalarValue]
values) \(Int
idx, ScalarValue
val) ->
        case (ScalarType
resultType, ScalarValue
val) of
            (ScalarType
_, ScalarValue
ScalarNull) ->
                Ptr Word64 -> Int -> IO ()
markInvalid Ptr Word64
validityPtr Int
idx
            (ScalarType
ScalarTypeBoolean, ScalarBoolean Bool
flag) -> do
                Ptr Word64 -> Int -> IO ()
markValid Ptr Word64
validityPtr Int
idx
                Ptr Word8 -> Int -> Word8 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff (Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8) Int
idx (if Bool
flag then Word8
1 else Word8
0)
            (ScalarType
ScalarTypeBigInt, ScalarInteger Int64
intval) -> do
                Ptr Word64 -> Int -> IO ()
markValid Ptr Word64
validityPtr Int
idx
                Ptr Int64 -> Int -> Int64 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff (Ptr () -> Ptr Int64
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Int64) Int
idx Int64
intval
            (ScalarType
ScalarTypeDouble, ScalarDouble Double
dbl) -> do
                Ptr Word64 -> Int -> IO ()
markValid Ptr Word64
validityPtr Int
idx
                Ptr Double -> Int -> Double -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff (Ptr () -> Ptr Double
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Double) Int
idx Double
dbl
            (ScalarType
ScalarTypeVarchar, ScalarText Text
txt) -> do
                Ptr Word64 -> Int -> IO ()
markValid Ptr Word64
validityPtr Int
idx
                Text -> (CStringLen -> IO ()) -> IO ()
forall a. Text -> (CStringLen -> IO a) -> IO a
TextForeign.withCStringLen Text
txt \(CString
ptr, Int
len) ->
                    DuckDBVector -> Word64 -> CString -> Word64 -> IO ()
c_duckdb_vector_assign_string_element_len DuckDBVector
outVec (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx) CString
ptr (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
            (ScalarType, ScalarValue)
_ ->
                SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO ()) -> SQLError -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Text -> SQLError
functionInvocationError (Text -> SQLError) -> Text -> SQLError
forall a b. (a -> b) -> a -> b
$
                        String -> Text
Text.pack String
"duckdb-simple: result type mismatch when materialising scalar function output"

markInvalid :: Ptr Word64 -> Int -> IO ()
markInvalid :: Ptr Word64 -> Int -> IO ()
markInvalid Ptr Word64
validity Int
idx
    | Ptr Word64
validity Ptr Word64 -> Ptr Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word64
forall a. Ptr a
nullPtr = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    | Bool
otherwise = Ptr Word64 -> Word64 -> IO ()
c_duckdb_validity_set_row_invalid Ptr Word64
validity (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx)

markValid :: Ptr Word64 -> Int -> IO ()
markValid :: Ptr Word64 -> Int -> IO ()
markValid Ptr Word64
validity Int
idx
    | Ptr Word64
validity Ptr Word64 -> Ptr Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word64
forall a. Ptr a
nullPtr = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    | Bool
otherwise = Ptr Word64 -> Word64 -> IO ()
c_duckdb_validity_set_row_valid Ptr Word64
validity (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx)

isNullValue :: ScalarValue -> Bool
isNullValue :: ScalarValue -> Bool
isNullValue = \case
    ScalarValue
ScalarNull -> Bool
True
    ScalarValue
_ -> Bool
False

argumentConversionError :: Int -> [SomeException] -> SQLError
argumentConversionError :: Int -> [SomeException] -> SQLError
argumentConversionError Int
idx [SomeException]
err =
    let message :: Text
message =
            [Text] -> Text
Text.concat
                [ String -> Text
Text.pack String
"duckdb-simple: unable to convert argument #"
                , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
                , String -> Text
Text.pack String
": "
                , String -> Text
Text.pack ([SomeException] -> String
forall a. Show a => a -> String
show [SomeException]
err)
                ]
     in Text -> SQLError
functionInvocationError Text
message

functionInvocationError :: Text -> SQLError
functionInvocationError :: Text -> SQLError
functionInvocationError Text
message =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage = Text
message
        , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
forall a. Maybe a
Nothing
        , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Maybe Query
forall a. Maybe a
Nothing
        }

fetchResultError :: Ptr DuckDBResult -> IO Text
fetchResultError :: Ptr DuckDBResult -> IO Text
fetchResultError Ptr DuckDBResult
resPtr = do
    CString
msgPtr <- Ptr DuckDBResult -> IO CString
c_duckdb_result_error Ptr DuckDBResult
resPtr
    if CString
msgPtr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr
        then Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Text
Text.pack String
"duckdb-simple: DROP FUNCTION failed")
        else 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
msgPtr

qualifyIdentifier :: Text -> Text
qualifyIdentifier :: Text -> Text
qualifyIdentifier Text
rawName =
    let parts :: [Text]
parts = HasCallStack => Text -> Text -> [Text]
Text -> Text -> [Text]
Text.splitOn Text
"." Text
rawName
     in Text -> [Text] -> Text
Text.intercalate (String -> Text
Text.pack String
".") ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
quoteIdent [Text]
parts)

quoteIdent :: Text -> Text
quoteIdent :: Text -> Text
quoteIdent Text
ident =
    [Text] -> Text
Text.concat
        [ String -> Text
Text.pack String
"\""
        , HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace (String -> Text
Text.pack String
"\"") (String -> Text
Text.pack String
"\"\"") Text
ident
        , String -> Text
Text.pack String
"\""
        ]

foreign import ccall "wrapper"
    mkScalarFun :: (DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()) -> IO DuckDBScalarFunctionFun

foreign import ccall "wrapper"
    mkDeleteCallback :: (Ptr () -> IO ()) -> IO DuckDBDeleteCallback