duckdb-simple-0.1.1.2: src/Database/DuckDB/Simple/Function.hs
{-# 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 '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.Bits ((.|.), shiftL)
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Ratio ((%))
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 Data.Time.Calendar (Day, fromGregorian)
import Data.Time.Clock (UTCTime (..))
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
import Data.Time.LocalTime
( LocalTime (..)
, TimeOfDay (..)
, minutesToTimeZone
, utc
, utcToLocalTime
)
import Database.DuckDB.FFI
import Database.DuckDB.Simple.FromField
( BigNum (..)
, BitString (..)
, DecimalValue (..)
, Field (..)
, FieldValue (..)
, FromField (..)
, IntervalValue (..)
, TimeWithZone (..)
)
import Database.DuckDB.Simple.Internal
( Connection
, Query (..)
, SQLError (..)
, withConnectionHandle
, withQueryCString
)
import Foreign.C.String (peekCString)
import Foreign.C.Types (CBool (..), CDouble (..))
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 (peek, 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 _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Int16 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Int32 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Int64 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger value)
instance FunctionResult Word where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Word16 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Word32 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Word64 where
scalarReturnType _ = ScalarTypeBigInt
toScalarValue value = pure (ScalarInteger (fromIntegral value))
instance FunctionResult Double where
scalarReturnType _ = ScalarTypeDouble
toScalarValue value = pure (ScalarDouble value)
instance FunctionResult Float where
scalarReturnType _ = ScalarTypeDouble
toScalarValue value = pure (ScalarDouble (realToFrac value))
instance FunctionResult Bool where
scalarReturnType _ = ScalarTypeBoolean
toScalarValue value = pure (ScalarBoolean value)
instance FunctionResult Text where
scalarReturnType _ = ScalarTypeVarchar
toScalarValue value = pure (ScalarText value)
instance FunctionResult String where
scalarReturnType _ = ScalarTypeVarchar
toScalarValue value = pure (ScalarText (Text.pack value))
instance (FunctionResult a) => FunctionResult (Maybe a) where
scalarReturnType _ = scalarReturnType (Proxy :: Proxy a)
toScalarValue Nothing = pure ScalarNull
toScalarValue (Just value) = toScalarValue value
-- | Argument types supported by the scalar function machinery.
class FunctionArg a where
argumentType :: Proxy a -> DuckDBType
instance FunctionArg Int where
argumentType _ = DuckDBTypeBigInt
instance FunctionArg Int16 where
argumentType _ = DuckDBTypeSmallInt
instance FunctionArg Int32 where
argumentType _ = DuckDBTypeInteger
instance FunctionArg Int64 where
argumentType _ = DuckDBTypeBigInt
instance FunctionArg Word where
argumentType _ = DuckDBTypeUBigInt
instance FunctionArg Word16 where
argumentType _ = DuckDBTypeUSmallInt
instance FunctionArg Word32 where
argumentType _ = DuckDBTypeUInteger
instance FunctionArg Word64 where
argumentType _ = DuckDBTypeUBigInt
instance FunctionArg Double where
argumentType _ = DuckDBTypeDouble
instance FunctionArg Float where
argumentType _ = DuckDBTypeFloat
instance FunctionArg Bool where
argumentType _ = DuckDBTypeBoolean
instance FunctionArg Text where
argumentType _ = DuckDBTypeVarchar
instance FunctionArg String where
argumentType _ = DuckDBTypeVarchar
instance (FunctionArg a) => FunctionArg (Maybe a) where
argumentType _ = argumentType (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 _ = []
returnType _ = scalarReturnType (Proxy :: Proxy a)
isVolatile _ = False
applyFunction [] value = toScalarValue value
applyFunction _ _ = throwIO (functionInvocationError (Text.pack "unexpected arguments supplied"))
instance {-# OVERLAPPING #-} (FunctionResult a) => Function (IO a) where
argumentTypes _ = []
returnType _ = scalarReturnType (Proxy :: Proxy a)
isVolatile _ = True
applyFunction [] action = action >>= toScalarValue
applyFunction _ _ = throwIO (functionInvocationError (Text.pack "unexpected arguments supplied"))
instance {-# OVERLAPPABLE #-} (FromField a, FunctionArg a, Function r) => Function (a -> r) where
argumentTypes _ = argumentType (Proxy :: Proxy a) : argumentTypes (Proxy :: Proxy r)
returnType _ = returnType (Proxy :: Proxy r)
isVolatile _ = isVolatile (Proxy :: Proxy r)
applyFunction [] _ =
throwIO (functionInvocationError (Text.pack "insufficient arguments supplied"))
applyFunction (field : rest) fn =
case fromField field of
Errors err -> throwIO (argumentConversionError (fieldIndex field) err)
Ok value -> applyFunction rest (fn value)
-- | Register a Haskell function under the supplied name.
createFunction :: forall f. (Function f) => Connection -> Text -> f -> IO ()
createFunction conn name fn = do
funPtr <- mkScalarFun (scalarFunctionHandler fn)
funPtrStable <- newStablePtr funPtr
destroyCb <- mkDeleteCallback releaseFunctionPtr
let release = destroyFunctionResources funPtr funPtrStable destroyCb
bracket c_duckdb_create_scalar_function cleanupScalarFunction \scalarFun ->
(`onException` release) $ do
TextForeign.withCString name $ \cName ->
c_duckdb_scalar_function_set_name scalarFun cName
forM_ (argumentTypes (Proxy :: Proxy f)) \dtype ->
withLogicalType dtype $ \logical ->
c_duckdb_scalar_function_add_parameter scalarFun logical
withLogicalType (duckTypeForScalar (returnType (Proxy :: Proxy f))) $ \logical ->
c_duckdb_scalar_function_set_return_type scalarFun logical
when (isVolatile (Proxy :: Proxy f)) $
c_duckdb_scalar_function_set_volatile scalarFun
c_duckdb_scalar_function_set_function scalarFun funPtr
c_duckdb_scalar_function_set_extra_info scalarFun (castStablePtrToPtr funPtrStable) destroyCb
withConnectionHandle conn \connPtr -> do
rc <- c_duckdb_register_scalar_function connPtr 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 conn name =
do
outcome <-
try $
withConnectionHandle conn \connPtr -> do
let dropQuery =
Query $
Text.concat
[ Text.pack "DROP FUNCTION IF EXISTS "
, qualifyIdentifier name
]
withQueryCString dropQuery \sql ->
alloca \resPtr -> do
rc <- c_duckdb_query connPtr sql 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 () -> pure ()
Left 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.isInfixOf (Text.pack "Cannot drop internal catalog entry") (sqlErrorMessage err) -> return ()
| otherwise -> throwIO err
cleanupScalarFunction :: DuckDBScalarFunction -> IO ()
cleanupScalarFunction scalarFun =
alloca \ptr -> do
poke ptr scalarFun
c_duckdb_destroy_scalar_function ptr
destroyFunctionResources :: DuckDBScalarFunctionFun -> StablePtr DuckDBScalarFunctionFun -> DuckDBDeleteCallback -> IO ()
destroyFunctionResources funPtr funPtrStable destroyCb = do
freeHaskellFunPtr funPtr
freeStablePtr funPtrStable
freeHaskellFunPtr destroyCb
releaseFunctionPtr :: Ptr () -> IO ()
releaseFunctionPtr rawPtr =
when (rawPtr /= nullPtr) $ do
let stablePtr = castPtrToStablePtr rawPtr :: StablePtr DuckDBScalarFunctionFun
funPtr <- deRefStablePtr stablePtr
freeHaskellFunPtr funPtr
freeStablePtr stablePtr
withLogicalType :: DuckDBType -> (DuckDBLogicalType -> IO a) -> IO a
withLogicalType dtype =
bracket
( do
logical <- c_duckdb_create_logical_type dtype
when (logical == nullPtr) $
throwIO $
functionInvocationError (Text.pack "duckdb-simple: failed to allocate logical type")
pure logical
)
destroyLogicalType
destroyLogicalType :: DuckDBLogicalType -> IO ()
destroyLogicalType logicalType =
alloca \ptr -> do
poke ptr logicalType
c_duckdb_destroy_logical_type ptr
duckTypeForScalar :: ScalarType -> DuckDBType
duckTypeForScalar = \case
ScalarTypeBoolean -> DuckDBTypeBoolean
ScalarTypeBigInt -> DuckDBTypeBigInt
ScalarTypeDouble -> DuckDBTypeDouble
ScalarTypeVarchar -> DuckDBTypeVarchar
scalarFunctionHandler :: forall f. (Function f) => f -> DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()
scalarFunctionHandler fn info chunk outVec = do
result <-
try do
rawColumnCount <- c_duckdb_data_chunk_get_column_count chunk
let columnCount = fromIntegral rawColumnCount :: Int
expected = length (argumentTypes (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 = fromIntegral rawRowCount :: Int
readers <- mapM (makeColumnReader chunk) [0 .. expected - 1]
rows <-
forM [0 .. rowCount - 1] \row ->
forM readers \reader ->
reader (fromIntegral 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 (err :: SomeException) -> do
c_duckdb_data_chunk_set_size chunk 0
let message = Text.pack (displayException err)
TextForeign.withCString message $ \cMsg ->
c_duckdb_scalar_function_set_error info cMsg
Right () -> pure ()
type ColumnReader = DuckDBIdx -> IO Field
makeColumnReader :: DuckDBDataChunk -> Int -> IO ColumnReader
makeColumnReader chunk columnIndex = do
vector <- c_duckdb_data_chunk_get_vector chunk (fromIntegral columnIndex)
logical <- c_duckdb_vector_get_column_type vector
dtype <- c_duckdb_get_type_id logical
decimalInfo <-
if dtype == DuckDBTypeDecimal
then do
width <- c_duckdb_decimal_width logical
scale <- c_duckdb_decimal_scale logical
pure (Just (width, scale))
else pure Nothing
enumInternal <-
if dtype == DuckDBTypeEnum
then Just <$> c_duckdb_enum_internal_type logical
else pure Nothing
destroyLogicalType logical
dataPtr <- c_duckdb_vector_get_data vector
validity <- c_duckdb_vector_get_validity vector
let name = Text.pack ("arg" <> show columnIndex)
pure \rowIdx -> do
valid <- isRowValid validity rowIdx
value <-
if valid
then fetchValue dtype decimalInfo enumInternal dataPtr rowIdx
else pure FieldNull
pure
Field
{ fieldName = name
, fieldIndex = columnIndex
, fieldValue = value
}
isRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool
isRowValid validity rowIdx
| validity == nullPtr = pure True
| otherwise = do
CBool flag <- c_duckdb_validity_row_is_valid validity rowIdx
pure (flag /= 0)
fetchValue :: DuckDBType -> Maybe (Word8, Word8) -> Maybe DuckDBType -> Ptr () -> DuckDBIdx -> IO FieldValue
fetchValue dtype decimalInfo enumInternal dataPtr rowIdx =
let idx = fromIntegral rowIdx
in case dtype of
DuckDBTypeBoolean -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word8) idx
pure (FieldBool (value /= 0))
DuckDBTypeTinyInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Int8) idx
pure (FieldInt8 value)
DuckDBTypeSmallInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Int16) idx
pure (FieldInt16 value)
DuckDBTypeInteger -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Int32) idx
pure (FieldInt32 value)
DuckDBTypeBigInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Int64) idx
pure (FieldInt64 value)
DuckDBTypeHugeInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBHugeInt) idx
pure (FieldHugeInt (duckDBHugeIntToInteger value))
DuckDBTypeUTinyInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word8) idx
pure (FieldWord8 value)
DuckDBTypeUSmallInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word16) idx
pure (FieldWord16 value)
DuckDBTypeUInteger -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word32) idx
pure (FieldWord32 value)
DuckDBTypeUBigInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word64) idx
pure (FieldWord64 value)
DuckDBTypeUHugeInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBUHugeInt) idx
pure (FieldUHugeInt (duckDBUHugeIntToInteger value))
DuckDBTypeFloat -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Float) idx
pure (FieldFloat value)
DuckDBTypeDouble -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Double) idx
pure (FieldDouble value)
DuckDBTypeVarchar -> FieldText <$> decodeDuckText dataPtr rowIdx
DuckDBTypeUUID -> FieldText <$> decodeDuckText dataPtr rowIdx
DuckDBTypeBlob -> FieldBlob <$> decodeDuckBlob dataPtr rowIdx
DuckDBTypeDate -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBDate) idx
FieldDate <$> decodeDuckDBDate value
DuckDBTypeTime -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTime) idx
FieldTime <$> decodeDuckDBTime value
DuckDBTypeTimeNs -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimeNs) idx
pure (FieldTime (decodeDuckDBTimeNs value))
DuckDBTypeTimeTz -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimeTz) idx
FieldTimeTZ <$> decodeDuckDBTimeTz value
DuckDBTypeTimestamp -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestamp) idx
FieldTimestamp <$> decodeDuckDBTimestamp value
DuckDBTypeTimestampS -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampS) idx
FieldTimestamp <$> decodeDuckDBTimestampSeconds value
DuckDBTypeTimestampMs -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampMs) idx
FieldTimestamp <$> decodeDuckDBTimestampMilliseconds value
DuckDBTypeTimestampNs -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestampNs) idx
FieldTimestamp <$> decodeDuckDBTimestampNanoseconds value
DuckDBTypeTimestampTz -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBTimestamp) idx
FieldTimestampTZ <$> decodeDuckDBTimestampUTCTime value
DuckDBTypeInterval -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBInterval) idx
pure (FieldInterval (intervalValueFromDuckDB value))
DuckDBTypeDecimal -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBDecimal) idx
case decimalInfo of
Just (width, scale) -> do
decimal <- decimalValueFromDuckDB width scale value
pure (FieldDecimal decimal)
Nothing ->
throwIO $
functionInvocationError $
Text.pack "duckdb-simple: missing decimal metadata for scalar function argument"
DuckDBTypeBit -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBBit) idx
FieldBit <$> decodeDuckDBBit value
DuckDBTypeBigNum -> do
value <- peekElemOff (castPtr dataPtr :: Ptr DuckDBBignum) idx
FieldBigNum <$> decodeDuckDBBigNum value
DuckDBTypeEnum -> do
case enumInternal of
Just DuckDBTypeUTinyInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word8) idx
pure (FieldEnum (fromIntegral value))
Just DuckDBTypeUSmallInt -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word16) idx
pure (FieldEnum (fromIntegral value))
Just DuckDBTypeUInteger -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Word32) idx
pure (FieldEnum value)
_ ->
throwIO $
functionInvocationError $
Text.pack "duckdb-simple: unsupported enum internal storage type for scalar function argument"
DuckDBTypeSQLNull ->
pure FieldNull
DuckDBTypeStringLiteral -> FieldText <$> decodeDuckText dataPtr rowIdx
DuckDBTypeIntegerLiteral -> do
value <- peekElemOff (castPtr dataPtr :: Ptr Int64) idx
pure (FieldInt64 value)
other ->
throwIO $
functionInvocationError $
Text.concat
[ Text.pack "duckdb-simple: unsupported argument type "
, Text.pack (show other)
]
decodeDuckText :: Ptr () -> DuckDBIdx -> IO Text
decodeDuckText dataPtr rowIdx = do
let bytePtr = castPtr dataPtr :: Ptr Word8
offset = fromIntegral rowIdx * duckdbStringTSize
stringPtr = castPtr (bytePtr `plusPtr` offset) :: Ptr DuckDBStringT
len <- c_duckdb_string_t_length stringPtr
if len == 0
then pure Text.empty
else do
cstr <- c_duckdb_string_t_data stringPtr
bytes <- BS.packCStringLen (cstr, fromIntegral len)
pure (TE.decodeUtf8 bytes)
decodeDuckBlob :: Ptr () -> DuckDBIdx -> IO BS.ByteString
decodeDuckBlob dataPtr rowIdx = do
let bytePtr = castPtr dataPtr :: Ptr Word8
offset = fromIntegral rowIdx * duckdbStringTSize
stringPtr = castPtr (bytePtr `plusPtr` offset) :: Ptr DuckDBStringT
len <- c_duckdb_string_t_length stringPtr
if len == 0
then pure BS.empty
else do
ptr <- c_duckdb_string_t_data stringPtr
BS.packCStringLen (ptr, fromIntegral len)
duckdbStringTSize :: Int
duckdbStringTSize = 16
decodeDuckDBDate :: DuckDBDate -> IO Day
decodeDuckDBDate raw =
alloca \ptr -> do
c_duckdb_from_date raw ptr
dateStruct <- peek ptr
pure (dateStructToDay dateStruct)
decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay
decodeDuckDBTime raw =
alloca \ptr -> do
c_duckdb_from_time raw ptr
timeStruct <- peek ptr
pure (timeStructToTimeOfDay timeStruct)
decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTime
decodeDuckDBTimestamp raw =
alloca \ptr -> do
c_duckdb_from_timestamp raw ptr
DuckDBTimestampStruct{duckDBTimestampStructDate = dateStruct, duckDBTimestampStructTime = timeStruct} <- peek ptr
pure
LocalTime
{ localDay = dateStructToDay dateStruct
, localTimeOfDay = timeStructToTimeOfDay timeStruct
}
dateStructToDay :: DuckDBDateStruct -> Day
dateStructToDay DuckDBDateStruct{duckDBDateStructYear, duckDBDateStructMonth, duckDBDateStructDay} =
fromGregorian (fromIntegral duckDBDateStructYear) (fromIntegral duckDBDateStructMonth) (fromIntegral duckDBDateStructDay)
timeStructToTimeOfDay :: DuckDBTimeStruct -> TimeOfDay
timeStructToTimeOfDay DuckDBTimeStruct{duckDBTimeStructHour, duckDBTimeStructMinute, duckDBTimeStructSecond, duckDBTimeStructMicros} =
let secondsInt = fromIntegral duckDBTimeStructSecond :: Integer
micros = fromIntegral duckDBTimeStructMicros :: Integer
fractional = fromRational (micros % 1000000)
totalSeconds = fromInteger secondsInt + fractional
in TimeOfDay
(fromIntegral duckDBTimeStructHour)
(fromIntegral duckDBTimeStructMinute)
totalSeconds
decodeDuckDBTimeNs :: DuckDBTimeNs -> TimeOfDay
decodeDuckDBTimeNs (DuckDBTimeNs nanos) =
let (hours, remainderHours) = nanos `divMod` (60 * 60 * 1000000000)
(minutes, remainderMinutes) = remainderHours `divMod` (60 * 1000000000)
(seconds, fractionalNanos) = remainderMinutes `divMod` 1000000000
fractional = fromRational (toInteger fractionalNanos % 1000000000)
totalSeconds = fromIntegral seconds + fractional
in TimeOfDay
(fromIntegral hours)
(fromIntegral minutes)
totalSeconds
decodeDuckDBTimeTz :: DuckDBTimeTz -> IO TimeWithZone
decodeDuckDBTimeTz raw =
alloca \ptr -> do
c_duckdb_from_time_tz raw ptr
DuckDBTimeTzStruct{duckDBTimeTzStructTime = timeStruct, duckDBTimeTzStructOffset = offset} <- peek ptr
let timeOfDay = timeStructToTimeOfDay timeStruct
minutes = fromIntegral offset `div` 60
zone = minutesToTimeZone minutes
pure TimeWithZone{timeWithZoneTime = timeOfDay, timeWithZoneZone = zone}
decodeDuckDBTimestampSeconds :: DuckDBTimestampS -> IO LocalTime
decodeDuckDBTimestampSeconds (DuckDBTimestampS seconds) =
decodeDuckDBTimestamp (DuckDBTimestamp (seconds * 1000000))
decodeDuckDBTimestampMilliseconds :: DuckDBTimestampMs -> IO LocalTime
decodeDuckDBTimestampMilliseconds (DuckDBTimestampMs millis) =
decodeDuckDBTimestamp (DuckDBTimestamp (millis * 1000))
decodeDuckDBTimestampNanoseconds :: DuckDBTimestampNs -> IO LocalTime
decodeDuckDBTimestampNanoseconds (DuckDBTimestampNs nanos) = do
let utcTime = posixSecondsToUTCTime (fromRational (toInteger nanos % 1000000000))
pure (utcToLocalTime utc utcTime)
decodeDuckDBTimestampUTCTime :: DuckDBTimestamp -> IO UTCTime
decodeDuckDBTimestampUTCTime (DuckDBTimestamp micros) =
pure (posixSecondsToUTCTime (fromRational (toInteger micros % 1000000)))
intervalValueFromDuckDB :: DuckDBInterval -> IntervalValue
intervalValueFromDuckDB DuckDBInterval{duckDBIntervalMonths, duckDBIntervalDays, duckDBIntervalMicros} =
IntervalValue
{ intervalMonths = duckDBIntervalMonths
, intervalDays = duckDBIntervalDays
, intervalMicros = duckDBIntervalMicros
}
decimalValueFromDuckDB :: Word8 -> Word8 -> DuckDBDecimal -> IO DecimalValue
decimalValueFromDuckDB width scale rawDecimal =
alloca \ptr -> do
let decimal = rawDecimal{duckDBDecimalWidth = width, duckDBDecimalScale = scale}
poke ptr decimal
CDouble approx <- c_duckdb_decimal_to_double ptr
pure
DecimalValue
{ decimalWidth = width
, decimalScale = scale
, decimalInteger = duckDBHugeIntToInteger (duckDBDecimalValue decimal)
, decimalApproximate = realToFrac approx
}
duckDBHugeIntToInteger :: DuckDBHugeInt -> Integer
duckDBHugeIntToInteger DuckDBHugeInt{duckDBHugeIntLower, duckDBHugeIntUpper} =
(fromIntegral duckDBHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBHugeIntLower
duckDBUHugeIntToInteger :: DuckDBUHugeInt -> Integer
duckDBUHugeIntToInteger DuckDBUHugeInt{duckDBUHugeIntLower, duckDBUHugeIntUpper} =
(fromIntegral duckDBUHugeIntUpper `shiftL` 64) .|. fromIntegral duckDBUHugeIntLower
decodeDuckDBBit :: DuckDBBit -> IO BitString
decodeDuckDBBit DuckDBBit{duckDBBitData, duckDBBitSize}
| duckDBBitData == nullPtr || duckDBBitSize == 0 =
pure BitString{bitStringLength = 0, bitStringBytes = BS.empty}
| otherwise = do
let bitLength = duckDBBitSize
byteLength = fromIntegral ((bitLength + 7) `div` 8) :: Int
bytes <- BS.packCStringLen (castPtr duckDBBitData, byteLength)
pure BitString{bitStringLength = bitLength, bitStringBytes = bytes}
decodeDuckDBBigNum :: DuckDBBignum -> IO BigNum
decodeDuckDBBigNum DuckDBBignum{duckDBBignumData, duckDBBignumSize, duckDBBignumIsNegative}
| duckDBBignumData == nullPtr || duckDBBignumSize == 0 =
pure BigNum{bigNumIsNegative = duckDBBignumIsNegative /= 0, bigNumMagnitude = BS.empty}
| otherwise = do
let byteLength = fromIntegral duckDBBignumSize :: Int
bytes <- BS.packCStringLen (castPtr duckDBBignumData, byteLength)
pure BigNum{bigNumIsNegative = duckDBBignumIsNegative /= 0, bigNumMagnitude = bytes}
writeResults :: ScalarType -> [ScalarValue] -> DuckDBVector -> IO ()
writeResults resultType values outVec = do
let hasNulls = any isNullValue values
when hasNulls $
c_duckdb_vector_ensure_validity_writable outVec
dataPtr <- c_duckdb_vector_get_data outVec
validityPtr <- c_duckdb_vector_get_validity outVec
forM_ (zip [0 ..] values) \(idx, val) ->
case (resultType, val) of
(_, ScalarNull) ->
markInvalid validityPtr idx
(ScalarTypeBoolean, ScalarBoolean flag) -> do
markValid validityPtr idx
pokeElemOff (castPtr dataPtr :: Ptr Word8) idx (if flag then 1 else 0)
(ScalarTypeBigInt, ScalarInteger intval) -> do
markValid validityPtr idx
pokeElemOff (castPtr dataPtr :: Ptr Int64) idx intval
(ScalarTypeDouble, ScalarDouble dbl) -> do
markValid validityPtr idx
pokeElemOff (castPtr dataPtr :: Ptr Double) idx dbl
(ScalarTypeVarchar, ScalarText txt) -> do
markValid validityPtr idx
TextForeign.withCStringLen txt \(ptr, len) ->
c_duckdb_vector_assign_string_element_len outVec (fromIntegral idx) ptr (fromIntegral len)
_ ->
throwIO $
functionInvocationError $
Text.pack "duckdb-simple: result type mismatch when materialising scalar function output"
markInvalid :: Ptr Word64 -> Int -> IO ()
markInvalid validity idx
| validity == nullPtr = pure ()
| otherwise = c_duckdb_validity_set_row_invalid validity (fromIntegral idx)
markValid :: Ptr Word64 -> Int -> IO ()
markValid validity idx
| validity == nullPtr = pure ()
| otherwise = c_duckdb_validity_set_row_valid validity (fromIntegral idx)
isNullValue :: ScalarValue -> Bool
isNullValue = \case
ScalarNull -> True
_ -> False
argumentConversionError :: Int -> [SomeException] -> SQLError
argumentConversionError idx err =
let message =
Text.concat
[ Text.pack "duckdb-simple: unable to convert argument #"
, Text.pack (show (idx + 1))
, Text.pack ": "
, Text.pack (show err)
]
in functionInvocationError message
functionInvocationError :: Text -> SQLError
functionInvocationError message =
SQLError
{ sqlErrorMessage = message
, sqlErrorType = Nothing
, sqlErrorQuery = Nothing
}
fetchResultError :: Ptr DuckDBResult -> IO Text
fetchResultError resPtr = do
msgPtr <- c_duckdb_result_error resPtr
if msgPtr == nullPtr
then pure (Text.pack "duckdb-simple: DROP FUNCTION failed")
else Text.pack <$> peekCString msgPtr
qualifyIdentifier :: Text -> Text
qualifyIdentifier rawName =
let parts = Text.splitOn "." rawName
in Text.intercalate (Text.pack ".") (map quoteIdent parts)
quoteIdent :: Text -> Text
quoteIdent ident =
Text.concat
[ Text.pack "\""
, Text.replace (Text.pack "\"") (Text.pack "\"\"") ident
, Text.pack "\""
]
foreign import ccall "wrapper"
mkScalarFun :: (DuckDBFunctionInfo -> DuckDBDataChunk -> DuckDBVector -> IO ()) -> IO DuckDBScalarFunctionFun
foreign import ccall "wrapper"
mkDeleteCallback :: (Ptr () -> IO ()) -> IO DuckDBDeleteCallback