{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# 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 field-decoding and result-marshalling machinery for arguments and
return values.
-}
module Database.DuckDB.Simple.Function (
    ScalarType,
    ScalarValue,
    FunctionArg (),
    FunctionResult (),
    Function (..),
    createFunction,
    createFunctionWithState,
    deleteFunction,
) where

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

data ScalarFunctionResources = ScalarFunctionResources
    { ScalarFunctionResources -> DuckDBScalarFunctionFun
scalarFunctionExecPtr :: !DuckDBScalarFunctionFun
    , ScalarFunctionResources -> Maybe DuckDBScalarFunctionInitFun
scalarFunctionInitPtr :: !(Maybe DuckDBScalarFunctionInitFun)
    }

-- | 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. (HasCallStack, 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. (HasCallStack, 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. (HasCallStack, 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. (HasCallStack, 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
    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)
    resources <- newStablePtr ScalarFunctionResources{scalarFunctionExecPtr = funPtr, scalarFunctionInitPtr = Nothing}
    destroyCb <- mkDeleteCallback releaseFunctionResources
    let release = DuckDBScalarFunctionFun
-> Maybe DuckDBScalarFunctionInitFun
-> StablePtr ScalarFunctionResources
-> DuckDBDeleteCallback
-> IO ()
destroyRegistrationResources DuckDBScalarFunctionFun
funPtr Maybe DuckDBScalarFunctionInitFun
forall a. Maybe a
Nothing StablePtr ScalarFunctionResources
resources DuckDBDeleteCallback
destroyCb
    bracket c_duckdb_create_scalar_function 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 ScalarFunctionResources -> Ptr ()
forall a. StablePtr a -> Ptr ()
castStablePtrToPtr StablePtr ScalarFunctionResources
resources) DuckDBDeleteCallback
destroyCb
            Connection -> (DuckDBConnection -> IO ()) -> IO ()
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr -> do
                rc <- DuckDBConnection -> DuckDBScalarFunction -> IO DuckDBState
c_duckdb_register_scalar_function DuckDBConnection
connPtr DuckDBScalarFunction
scalarFun
                if rc == DuckDBSuccess
                    then pure ()
                    else throwIO (functionInvocationError (Text.pack "duckdb-simple: registering function failed"))

-- | Register a scalar function with per-worker thread-local state.
createFunctionWithState :: forall s f. (Function f) => Connection -> Text -> IO s -> (s -> f) -> IO ()
createFunctionWithState :: forall s f.
Function f =>
Connection -> Text -> IO s -> (s -> f) -> IO ()
createFunctionWithState Connection
conn Text
name IO s
initState s -> f
mkFn = do
    stateDestroyCb <- (Ptr () -> IO ()) -> IO DuckDBDeleteCallback
mkDeleteCallback Ptr () -> IO ()
releaseStablePtrData
    execPtr <- mkScalarFun (scalarFunctionHandlerWithState mkFn)
    initPtr <- mkScalarInitFun (scalarFunctionInitHandler stateDestroyCb initState)
    resources <-
        newStablePtr
            ScalarFunctionResources
                { scalarFunctionExecPtr = execPtr
                , scalarFunctionInitPtr = Just initPtr
                }
    destroyCb <- mkDeleteCallback releaseFunctionResources
    let release = DuckDBScalarFunctionFun
-> Maybe DuckDBScalarFunctionInitFun
-> StablePtr ScalarFunctionResources
-> DuckDBDeleteCallback
-> IO ()
destroyRegistrationResources DuckDBScalarFunctionFun
execPtr (DuckDBScalarFunctionInitFun -> Maybe DuckDBScalarFunctionInitFun
forall a. a -> Maybe a
Just DuckDBScalarFunctionInitFun
initPtr) StablePtr ScalarFunctionResources
resources DuckDBDeleteCallback
destroyCb 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
stateDestroyCb
    bracket c_duckdb_create_scalar_function 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
execPtr
            DuckDBScalarFunction -> DuckDBScalarFunctionInitFun -> IO ()
c_duckdb_scalar_function_set_init DuckDBScalarFunction
scalarFun DuckDBScalarFunctionInitFun
initPtr
            DuckDBScalarFunction -> Ptr () -> DuckDBDeleteCallback -> IO ()
c_duckdb_scalar_function_set_extra_info DuckDBScalarFunction
scalarFun (StablePtr ScalarFunctionResources -> Ptr ()
forall a. StablePtr a -> Ptr ()
castStablePtrToPtr StablePtr ScalarFunctionResources
resources) DuckDBDeleteCallback
destroyCb
            Connection -> (DuckDBConnection -> IO ()) -> IO ()
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr -> do
                rc <- DuckDBConnection -> DuckDBScalarFunction -> IO DuckDBState
c_duckdb_register_scalar_function DuckDBConnection
connPtr DuckDBScalarFunction
scalarFun
                if rc == DuckDBSuccess
                    then pure ()
                    else throwIO (functionInvocationError (Text.pack "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
        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
                            rc <- DuckDBConnection -> CString -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_query DuckDBConnection
connPtr CString
sql Ptr DuckDBResult
resPtr
                            if rc == DuckDBSuccess
                                then c_duckdb_destroy_result resPtr
                                else do
                                    errMsg <- fetchResultError resPtr
                                    c_duckdb_destroy_result resPtr
                                    throwIO
                                        SQLError
                                            { sqlErrorMessage = errMsg
                                            , sqlErrorType = Nothing
                                            , sqlErrorQuery = Just dropQuery
                                            }
        case 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. (HasCallStack, 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

destroyRegistrationResources ::
    DuckDBScalarFunctionFun ->
    Maybe DuckDBScalarFunctionInitFun ->
    StablePtr ScalarFunctionResources ->
    DuckDBDeleteCallback ->
    IO ()
destroyRegistrationResources :: DuckDBScalarFunctionFun
-> Maybe DuckDBScalarFunctionInitFun
-> StablePtr ScalarFunctionResources
-> DuckDBDeleteCallback
-> IO ()
destroyRegistrationResources DuckDBScalarFunctionFun
funPtr Maybe DuckDBScalarFunctionInitFun
mInitPtr StablePtr ScalarFunctionResources
resources DuckDBDeleteCallback
destroyCb = do
    DuckDBScalarFunctionFun -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBScalarFunctionFun
funPtr
    Maybe DuckDBScalarFunctionInitFun
-> (DuckDBScalarFunctionInitFun -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe DuckDBScalarFunctionInitFun
mInitPtr DuckDBScalarFunctionInitFun -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr
    StablePtr ScalarFunctionResources -> IO ()
forall a. StablePtr a -> IO ()
freeStablePtr StablePtr ScalarFunctionResources
resources
    DuckDBDeleteCallback -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr DuckDBDeleteCallback
destroyCb

releaseFunctionResources :: Ptr () -> IO ()
releaseFunctionResources :: Ptr () -> IO ()
releaseFunctionResources 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 ScalarFunctionResources
stablePtr = Ptr () -> StablePtr ScalarFunctionResources
forall a. Ptr () -> StablePtr a
castPtrToStablePtr Ptr ()
rawPtr :: StablePtr ScalarFunctionResources
        ScalarFunctionResources{scalarFunctionExecPtr, scalarFunctionInitPtr} <- StablePtr ScalarFunctionResources -> IO ScalarFunctionResources
forall a. StablePtr a -> IO a
deRefStablePtr StablePtr ScalarFunctionResources
stablePtr
        freeHaskellFunPtr scalarFunctionExecPtr
        forM_ scalarFunctionInitPtr freeHaskellFunPtr
        freeStablePtr 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
            logical <- DuckDBType -> IO DuckDBLogicalType
c_duckdb_create_logical_type DuckDBType
dtype
            when (logical == nullPtr) $
                throwIO $
                    functionInvocationError (Text.pack "duckdb-simple: failed to allocate logical type")
            pure logical
        )
        DuckDBLogicalType -> IO ()
destroyLogicalType

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
    result <-
        IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try do
            rawColumnCount <- DuckDBDataChunk -> IO Word64
c_duckdb_data_chunk_get_column_count DuckDBDataChunk
chunk
            let columnCount = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rawColumnCount :: 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))
            when (columnCount /= expected) $
                throwIO $
                    functionInvocationError $
                        Text.concat
                            [ Text.pack "duckdb-simple: function expected "
                            , Text.pack (show expected)
                            , Text.pack " arguments but received "
                            , Text.pack (show columnCount)
                            ]
            rawRowCount <- c_duckdb_data_chunk_get_size chunk
            let rowCount = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rawRowCount :: Int
            readers <- mapM (makeColumnReader chunk) [0 .. expected - 1]
            rows <-
                forM [0 .. rowCount - 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)
            results <- mapM (`applyFunction` fn) rows
            writeResults (returnType (Proxy :: Proxy f)) results outVec
            c_duckdb_data_chunk_set_size chunk (fromIntegral rowCount)
    case 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 ()

scalarFunctionHandlerWithState :: forall s f. (Function f) => (s -> f) -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandlerWithState :: forall s f.
Function f =>
(s -> f)
-> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandlerWithState s -> f
mkFn DuckDBFunctionInfo
info DuckDBDataChunk
chunk DuckDBVector
outVec = do
    statePtr <- DuckDBFunctionInfo -> IO (Ptr ())
c_duckdb_scalar_function_get_state DuckDBFunctionInfo
info
    if statePtr == nullPtr
        then
            TextForeign.withCString (Text.pack "duckdb-simple: scalar function state was not initialised") $
                c_duckdb_scalar_function_set_error info
        else do
            state <- deRefStablePtr (castPtrToStablePtr statePtr :: StablePtr s)
            scalarFunctionHandler (mkFn state) info chunk outVec

scalarFunctionInitHandler :: forall s. DuckDBDeleteCallback -> IO s -> DuckDBInitInfo -> IO ()
scalarFunctionInitHandler :: forall s. DuckDBDeleteCallback -> IO s -> DuckDBInitInfo -> IO ()
scalarFunctionInitHandler DuckDBDeleteCallback
destroyCb IO s
initState DuckDBInitInfo
info = do
    outcome <- IO s -> IO (Either SomeException s)
forall e a. Exception e => IO a -> IO (Either e a)
try IO s
initState
    case outcome of
        Left (SomeException
err :: SomeException) ->
            Text -> (CString -> IO ()) -> IO ()
forall a. Text -> (CString -> IO a) -> IO a
TextForeign.withCString (String -> Text
Text.pack (SomeException -> String
forall e. Exception e => e -> String
displayException SomeException
err)) ((CString -> IO ()) -> IO ()) -> (CString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
                DuckDBInitInfo -> CString -> IO ()
c_duckdb_scalar_function_init_set_error DuckDBInitInfo
info
        Right s
state -> do
            stable <- s -> IO (StablePtr s)
forall a. a -> IO (StablePtr a)
newStablePtr s
state
            c_duckdb_scalar_function_init_set_state info (castStablePtrToPtr stable) destroyCb

type ColumnReader = DuckDBIdx -> IO Field

makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader DuckDBDataChunk
chunk Int
columnIndex = do
    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)
    logical <- c_duckdb_vector_get_column_type vector
    dtype <- c_duckdb_get_type_id logical
    destroyLogicalType logical
    dataPtr <- c_duckdb_vector_get_data vector
    validity <- c_duckdb_vector_get_validity vector
    let 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)
    pure \Word64
rowIdx -> do
        value <- DuckDBType
-> DuckDBVector -> Ptr () -> Ptr Word64 -> Int -> IO FieldValue
materializeValue DuckDBType
dtype DuckDBVector
vector Ptr ()
dataPtr Ptr Word64
validity (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
rowIdx)
        pure
            Field
                { fieldName = name
                , fieldIndex = columnIndex
                , fieldValue = value
                }
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
    dataPtr <- DuckDBVector -> IO (Ptr ())
c_duckdb_vector_get_data DuckDBVector
outVec
    validityPtr <- c_duckdb_vector_get_validity outVec
    forM_ (zip [0 ..] 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. (HasCallStack, 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
    msgPtr <- Ptr DuckDBResult -> IO CString
c_duckdb_result_error Ptr DuckDBResult
resPtr
    if msgPtr == nullPtr
        then pure (Text.pack "duckdb-simple: DROP FUNCTION failed")
        else Text.pack <$> peekCString 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"
    mkScalarInitFun :: (DuckDBInitInfo -> IO ()) -> IO DuckDBScalarFunctionInitFun