{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TupleSections #-}

{- |
Module      : Database.DuckDB.Simple
Description : High-level DuckDB API in the duckdb-simple style.

The API mirrors the ergonomics of @sqlite-simple@ while being backed by the
DuckDB C API. It supports connection management, parameter binding, execution,
and typed result decoding. See 'README.md' for usage examples.
-}
module Database.DuckDB.Simple (
    -- * Connections
    Connection,
    open,
    close,
    withConnection,

    -- * Queries and statements
    Query (..),
    Statement,
    openStatement,
    closeStatement,
    withStatement,
    clearStatementBindings,
    namedParameterIndex,
    columnCount,
    columnName,
    executeStatement,
    execute,
    executeMany,
    execute_,
    bind,
    bindNamed,
    executeNamed,
    queryNamed,
    fold,
    fold_,
    foldNamed,
    withTransaction,
    query,
    queryWith,
    query_,
    queryWith_,
    nextRow,
    nextRowWith,

    -- * Errors and conversions
    SQLError (..),
    FormatError (..),
    ResultError (..),
    FieldParser,
    FromField (..),
    FromRow (..),
    RowParser,
    field,
    fieldWith,
    numFieldsRemaining,
    -- Re-export parameter helper types.
    ToField (..),
    ToRow (..),
    FieldBinding,
    NamedParam (..),
    Null (..),
    Only (..),
    (:.) (..),

    -- * User-defined scalar functions
    Function,
    createFunction,
    deleteFunction,
) where

import Control.Monad (forM, foldM, join, void, when, zipWithM, zipWithM_)
import Data.IORef (IORef, atomicModifyIORef', mkWeakIORef, newIORef, readIORef, writeIORef)
import Control.Exception (SomeException, bracket, finally, mask, onException, throwIO, try)
import qualified Data.ByteString as BS
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Maybe (isJust, isNothing)
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Foreign as TextForeign
import qualified Data.Text.Encoding as TextEncoding
import Data.Time.Calendar (Day, fromGregorian)
import Data.Time.LocalTime (LocalTime (..), TimeOfDay (..))
import Data.Word (Word16, Word32, Word64, Word8)
import Data.Ratio ((%))
import Database.DuckDB.FFI
import Database.DuckDB.Simple.FromField (
    Field (..),
    FieldValue (..),
    FieldParser,
    FromField (..),
    ResultError (..),
 )
import Database.DuckDB.Simple.FromRow (
    FromRow (..),
    RowParser,
    field,
    fieldWith,
    numFieldsRemaining,
    parseRow,
    rowErrorsToSqlError
    )
import Database.DuckDB.Simple.Function (Function, createFunction, deleteFunction)
import Database.DuckDB.Simple.Internal (
    Connection (..),
    ConnectionState (..),
    Query (..),
    SQLError (..),
    Statement (..),
    StatementState (..),
    StatementStream (..),
    StatementStreamChunk (..),
    StatementStreamChunkVector (..),
    StatementStreamColumn (..),
    StatementStreamState (..),
    withConnectionHandle,
    withQueryCString,
    withStatementHandle
    )
import Database.DuckDB.Simple.ToField (FieldBinding, NamedParam (..), ToField (..), bindFieldBinding, renderFieldBinding)
import Database.DuckDB.Simple.ToRow (ToRow (..))
import Database.DuckDB.Simple.Types (FormatError (..), Null (..), Only (..), (:.) (..))
import Foreign.C.String (CString, peekCString, withCString)
import Foreign.C.Types (CBool (..), CDouble (..))
import Foreign.Marshal.Alloc (alloca, free, malloc)
import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr)
import Foreign.Storable (peek, peekElemOff, poke)
import Database.DuckDB.Simple.Ok (Ok(..))

-- | Open a DuckDB database located at the supplied path.
open :: FilePath -> IO Connection
open :: String -> IO Connection
open String
path =
    ((forall a. IO a -> IO a) -> IO Connection) -> IO Connection
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask \forall a. IO a -> IO a
restore -> do
        DuckDBDatabase
db <- IO DuckDBDatabase -> IO DuckDBDatabase
forall a. IO a -> IO a
restore (String -> IO DuckDBDatabase
openDatabase String
path)
        DuckDBConnection
conn <-
            IO DuckDBConnection -> IO DuckDBConnection
forall a. IO a -> IO a
restore (DuckDBDatabase -> IO DuckDBConnection
connectDatabase DuckDBDatabase
db)
                IO DuckDBConnection -> IO () -> IO DuckDBConnection
forall a b. IO a -> IO b -> IO a
`onException` DuckDBDatabase -> IO ()
closeDatabaseHandle DuckDBDatabase
db
        DuckDBDatabase -> DuckDBConnection -> IO Connection
createConnection DuckDBDatabase
db DuckDBConnection
conn
            IO Connection -> IO () -> IO Connection
forall a b. IO a -> IO b -> IO a
`onException` do
                DuckDBConnection -> IO ()
closeConnectionHandle DuckDBConnection
conn
                DuckDBDatabase -> IO ()
closeDatabaseHandle DuckDBDatabase
db

-- | Close a connection.  The operation is idempotent.
close :: Connection -> IO ()
close :: Connection -> IO ()
close Connection{IORef ConnectionState
connectionState :: IORef ConnectionState
connectionState :: Connection -> IORef ConnectionState
connectionState} =
    IO (IO ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
        IORef ConnectionState
-> (ConnectionState -> (ConnectionState, IO ())) -> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef ConnectionState
connectionState \case
            ConnectionState
ConnectionClosed -> (ConnectionState
ConnectionClosed, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
            openState :: ConnectionState
openState@(ConnectionOpen{}) ->
                (ConnectionState
ConnectionClosed, ConnectionState -> IO ()
closeHandles ConnectionState
openState)

-- | Run an action with a freshly opened connection, closing it afterwards.
withConnection :: FilePath -> (Connection -> IO a) -> IO a
withConnection :: forall a. String -> (Connection -> IO a) -> IO a
withConnection String
path = IO Connection
-> (Connection -> IO ()) -> (Connection -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (String -> IO Connection
open String
path) Connection -> IO ()
close

-- | Prepare a SQL statement for execution.
openStatement :: Connection -> Query -> IO Statement
openStatement :: Connection -> Query -> IO Statement
openStatement Connection
conn Query
queryText =
    ((forall a. IO a -> IO a) -> IO Statement) -> IO Statement
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask \forall a. IO a -> IO a
restore -> do
        DuckDBPreparedStatement
handle <-
            IO DuckDBPreparedStatement -> IO DuckDBPreparedStatement
forall a. IO a -> IO a
restore (IO DuckDBPreparedStatement -> IO DuckDBPreparedStatement)
-> IO DuckDBPreparedStatement -> IO DuckDBPreparedStatement
forall a b. (a -> b) -> a -> b
$
                Connection
-> (DuckDBConnection -> IO DuckDBPreparedStatement)
-> IO DuckDBPreparedStatement
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr ->
                    Query
-> (CString -> IO DuckDBPreparedStatement)
-> IO DuckDBPreparedStatement
forall a. Query -> (CString -> IO a) -> IO a
withQueryCString Query
queryText \CString
sql ->
                        (Ptr DuckDBPreparedStatement -> IO DuckDBPreparedStatement)
-> IO DuckDBPreparedStatement
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBPreparedStatement
stmtPtr -> do
                            DuckDBState
rc <- DuckDBConnection
-> CString -> Ptr DuckDBPreparedStatement -> IO DuckDBState
c_duckdb_prepare DuckDBConnection
connPtr CString
sql Ptr DuckDBPreparedStatement
stmtPtr
                            DuckDBPreparedStatement
stmt <- Ptr DuckDBPreparedStatement -> IO DuckDBPreparedStatement
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBPreparedStatement
stmtPtr
                            if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                                then DuckDBPreparedStatement -> IO DuckDBPreparedStatement
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DuckDBPreparedStatement
stmt
                                else do
                                    Text
errMsg <- DuckDBPreparedStatement -> IO Text
fetchPrepareError DuckDBPreparedStatement
stmt
                                    Ptr DuckDBPreparedStatement -> IO ()
c_duckdb_destroy_prepare Ptr DuckDBPreparedStatement
stmtPtr
                                    SQLError -> IO DuckDBPreparedStatement
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO DuckDBPreparedStatement)
-> SQLError -> IO DuckDBPreparedStatement
forall a b. (a -> b) -> a -> b
$ Query -> Text -> SQLError
mkPrepareError Query
queryText Text
errMsg
        Connection -> DuckDBPreparedStatement -> Query -> IO Statement
createStatement Connection
conn DuckDBPreparedStatement
handle Query
queryText
            IO Statement -> IO () -> IO Statement
forall a b. IO a -> IO b -> IO a
`onException` DuckDBPreparedStatement -> IO ()
destroyPrepared DuckDBPreparedStatement
handle

-- | Finalise a prepared statement.  The operation is idempotent.
closeStatement :: Statement -> IO ()
closeStatement :: Statement -> IO ()
closeStatement stmt :: Statement
stmt@Statement{IORef StatementState
statementState :: IORef StatementState
statementState :: Statement -> IORef StatementState
statementState} = do
    Statement -> IO ()
resetStatementStream Statement
stmt
    IO (IO ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
        IORef StatementState
-> (StatementState -> (StatementState, IO ())) -> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef StatementState
statementState \case
            StatementState
StatementClosed -> (StatementState
StatementClosed, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
            StatementOpen{DuckDBPreparedStatement
statementHandle :: DuckDBPreparedStatement
statementHandle :: StatementState -> DuckDBPreparedStatement
statementHandle} ->
                (StatementState
StatementClosed, DuckDBPreparedStatement -> IO ()
destroyPrepared DuckDBPreparedStatement
statementHandle)

-- | Run an action with a prepared statement, closing it afterwards.
withStatement :: Connection -> Query -> (Statement -> IO a) -> IO a
withStatement :: forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
sql = IO Statement -> (Statement -> IO ()) -> (Statement -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Connection -> Query -> IO Statement
openStatement Connection
conn Query
sql) Statement -> IO ()
closeStatement

-- | Bind positional parameters to a prepared statement, replacing any previous bindings.
bind :: Statement -> [FieldBinding] -> IO ()
bind :: Statement -> [FieldBinding] -> IO ()
bind Statement
stmt [FieldBinding]
fields = do
    Statement -> IO ()
resetStatementStream Statement
stmt
    Statement -> (DuckDBPreparedStatement -> IO ()) -> IO ()
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
        let actual :: Int
actual = [FieldBinding] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [FieldBinding]
fields
        Int
expected <- (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBPreparedStatement -> IO DuckDBIdx
c_duckdb_nparams DuckDBPreparedStatement
handle)
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
actual Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
expected) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            Statement -> Text -> [FieldBinding] -> IO ()
forall a. Statement -> Text -> [FieldBinding] -> IO a
throwFormatErrorBindings Statement
stmt (Int -> Int -> Text
forall {a} {a}. (Show a, Show a) => a -> a -> Text
parameterCountMessage Int
expected Int
actual) [FieldBinding]
fields
        [Maybe Text]
parameterNames <- DuckDBPreparedStatement -> Int -> IO [Maybe Text]
fetchParameterNames DuckDBPreparedStatement
handle Int
expected
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Maybe Text -> Bool) -> [Maybe Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Maybe Text -> Bool
forall a. Maybe a -> Bool
isJust [Maybe Text]
parameterNames) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            Statement -> Text -> [FieldBinding] -> IO ()
forall a. Statement -> Text -> [FieldBinding] -> IO a
throwFormatErrorBindings Statement
stmt (String -> Text
Text.pack String
"duckdb-simple: statement defines named parameters; use executeNamed or bindNamed") [FieldBinding]
fields
    Statement -> IO ()
clearStatementBindings Statement
stmt
    (Int -> FieldBinding -> IO ()) -> [Int] -> [FieldBinding] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Int -> FieldBinding -> IO ()
apply [Int
1 ..] [FieldBinding]
fields
  where
    parameterCountMessage :: a -> a -> Text
parameterCountMessage a
expected a
actual =
        String -> Text
Text.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$
            String
"duckdb-simple: SQL query contains "
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
expected
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" parameter(s), but "
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
actual
                String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" argument(s) were supplied"

    apply :: Int -> FieldBinding -> IO ()
    apply :: Int -> FieldBinding -> IO ()
apply Int
idx = Statement -> DuckDBIdx -> FieldBinding -> IO ()
bindFieldBinding Statement
stmt (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx :: DuckDBIdx)

-- | Bind named parameters to a prepared statement, preserving any positional bindings.
bindNamed :: Statement -> [NamedParam] -> IO ()
bindNamed :: Statement -> [NamedParam] -> IO ()
bindNamed Statement
stmt [NamedParam]
params =
    let bindings :: [(Text, FieldBinding)]
bindings = (NamedParam -> (Text, FieldBinding))
-> [NamedParam] -> [(Text, FieldBinding)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Text
name := a
value) -> (Text
name, a -> FieldBinding
forall a. ToField a => a -> FieldBinding
toField a
value)) [NamedParam]
params
        parameterCountMessage :: a -> a -> Text
parameterCountMessage a
expected a
actual =
            String -> Text
Text.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$
                String
"duckdb-simple: SQL query contains "
                    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
expected
                    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" named parameter(s), but "
                    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
actual
                    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" argument(s) were supplied"
        unknownNameMessage :: Text -> Text
unknownNameMessage Text
name =
            [Text] -> Text
Text.concat
                [ String -> Text
Text.pack String
"duckdb-simple: unknown named parameter "
                , Text
name
                ]
        apply :: (Text, FieldBinding) -> IO ()
apply (Text
name, FieldBinding
binding) = do
            Maybe Int
mIdx <- Statement -> Text -> IO (Maybe Int)
namedParameterIndex Statement
stmt Text
name
            case Maybe Int
mIdx of
                Maybe Int
Nothing ->
                    Statement -> Text -> [(Text, FieldBinding)] -> IO ()
forall a. Statement -> Text -> [(Text, FieldBinding)] -> IO a
throwFormatErrorNamed Statement
stmt (Text -> Text
unknownNameMessage Text
name) [(Text, FieldBinding)]
bindings
                Just Int
idx -> Statement -> DuckDBIdx -> FieldBinding -> IO ()
bindFieldBinding Statement
stmt (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx :: DuckDBIdx) FieldBinding
binding
     in do
            Statement -> IO ()
resetStatementStream Statement
stmt
            Statement -> (DuckDBPreparedStatement -> IO ()) -> IO ()
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
                let actual :: Int
actual = [(Text, FieldBinding)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Text, FieldBinding)]
bindings
                Int
expected <- (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBPreparedStatement -> IO DuckDBIdx
c_duckdb_nparams DuckDBPreparedStatement
handle)
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
actual Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
expected) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Statement -> Text -> [(Text, FieldBinding)] -> IO ()
forall a. Statement -> Text -> [(Text, FieldBinding)] -> IO a
throwFormatErrorNamed Statement
stmt (Int -> Int -> Text
forall {a} {a}. (Show a, Show a) => a -> a -> Text
parameterCountMessage Int
expected Int
actual) [(Text, FieldBinding)]
bindings
                [Maybe Text]
parameterNames <- DuckDBPreparedStatement -> Int -> IO [Maybe Text]
fetchParameterNames DuckDBPreparedStatement
handle Int
expected
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Maybe Text -> Bool) -> [Maybe Text] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Maybe Text -> Bool
forall a. Maybe a -> Bool
isNothing [Maybe Text]
parameterNames Bool -> Bool -> Bool
&& Int
expected Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                    Statement -> Text -> [(Text, FieldBinding)] -> IO ()
forall a. Statement -> Text -> [(Text, FieldBinding)] -> IO a
throwFormatErrorNamed Statement
stmt (String -> Text
Text.pack String
"duckdb-simple: statement does not define named parameters; use positional bindings or adjust the SQL") [(Text, FieldBinding)]
bindings
            Statement -> IO ()
clearStatementBindings Statement
stmt
            ((Text, FieldBinding) -> IO ()) -> [(Text, FieldBinding)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Text, FieldBinding) -> IO ()
apply [(Text, FieldBinding)]
bindings

fetchParameterNames :: DuckDBPreparedStatement -> Int -> IO [Maybe Text]
fetchParameterNames :: DuckDBPreparedStatement -> Int -> IO [Maybe Text]
fetchParameterNames DuckDBPreparedStatement
handle Int
count =
    [Int] -> (Int -> IO (Maybe Text)) -> IO [Maybe Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
1 .. Int
count] \Int
idx -> do
        CString
namePtr <- DuckDBPreparedStatement -> DuckDBIdx -> IO CString
c_duckdb_parameter_name DuckDBPreparedStatement
handle (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx)
        if CString
namePtr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr
            then Maybe Text -> IO (Maybe Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Text
forall a. Maybe a
Nothing
            else do
                Text
name <- 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
namePtr
                Ptr () -> IO ()
c_duckdb_free (CString -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr CString
namePtr)
                let normalized :: Text
normalized = Text -> Text
normalizeName Text
name
                if Text
normalized Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
idx)
                    then Maybe Text -> IO (Maybe Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Text
forall a. Maybe a
Nothing
                    else Maybe Text -> IO (Maybe Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
name)

-- | Remove all parameter bindings associated with a prepared statement.
clearStatementBindings :: Statement -> IO ()
clearStatementBindings :: Statement -> IO ()
clearStatementBindings Statement
stmt =
    Statement -> (DuckDBPreparedStatement -> IO ()) -> IO ()
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
        DuckDBState
rc <- DuckDBPreparedStatement -> IO DuckDBState
c_duckdb_clear_bindings DuckDBPreparedStatement
handle
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
/= DuckDBState
DuckDBSuccess) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Text
err <- DuckDBPreparedStatement -> IO Text
fetchPrepareError DuckDBPreparedStatement
handle
            SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO ()) -> SQLError -> IO ()
forall a b. (a -> b) -> a -> b
$ Query -> Text -> SQLError
mkPrepareError (Statement -> Query
statementQuery Statement
stmt) Text
err

-- | Look up the 1-based index of a named placeholder.
namedParameterIndex :: Statement -> Text -> IO (Maybe Int)
namedParameterIndex :: Statement -> Text -> IO (Maybe Int)
namedParameterIndex Statement
stmt Text
name =
    Statement
-> (DuckDBPreparedStatement -> IO (Maybe Int)) -> IO (Maybe Int)
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle ->
        let normalized :: Text
normalized = Text -> Text
normalizeName Text
name
         in Text -> (CString -> IO (Maybe Int)) -> IO (Maybe Int)
forall a. Text -> (CString -> IO a) -> IO a
TextForeign.withCString Text
normalized \CString
cName ->
                (Ptr DuckDBIdx -> IO (Maybe Int)) -> IO (Maybe Int)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBIdx
idxPtr -> do
                    DuckDBState
rc <- DuckDBPreparedStatement
-> Ptr DuckDBIdx -> CString -> IO DuckDBState
c_duckdb_bind_parameter_index DuckDBPreparedStatement
handle Ptr DuckDBIdx
idxPtr CString
cName
                    if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                        then do
                            DuckDBIdx
idx <- Ptr DuckDBIdx -> IO DuckDBIdx
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBIdx
idxPtr
                            if DuckDBIdx
idx DuckDBIdx -> DuckDBIdx -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBIdx
0
                                then Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
                                else Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Maybe Int
forall a. a -> Maybe a
Just (DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
idx))
                        else Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing

-- | Retrieve the number of columns produced by the supplied prepared statement.
columnCount :: Statement -> IO Int
columnCount :: Statement -> IO Int
columnCount Statement
stmt =
    Statement -> (DuckDBPreparedStatement -> IO Int) -> IO Int
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt ((DuckDBPreparedStatement -> IO Int) -> IO Int)
-> (DuckDBPreparedStatement -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \DuckDBPreparedStatement
handle ->
        (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBPreparedStatement -> IO DuckDBIdx
c_duckdb_prepared_statement_column_count DuckDBPreparedStatement
handle)

-- | Look up the zero-based column name exposed by a prepared statement result.
columnName :: Statement -> Int -> IO Text
columnName :: Statement -> Int -> IO Text
columnName Statement
stmt Int
columnIndex
    | Int
columnIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = SQLError -> IO Text
forall e a. Exception e => e -> IO a
throwIO (Statement -> Int -> Maybe Int -> SQLError
columnIndexError Statement
stmt Int
columnIndex Maybe Int
forall a. Maybe a
Nothing)
    | Bool
otherwise =
        Statement -> (DuckDBPreparedStatement -> IO Text) -> IO Text
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
            Int
total <- (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBPreparedStatement -> IO DuckDBIdx
c_duckdb_prepared_statement_column_count DuckDBPreparedStatement
handle)
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
columnIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
total) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                SQLError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (Statement -> Int -> Maybe Int -> SQLError
columnIndexError Statement
stmt Int
columnIndex (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
total))
            CString
namePtr <- DuckDBPreparedStatement -> DuckDBIdx -> IO CString
c_duckdb_prepared_statement_column_name DuckDBPreparedStatement
handle (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
columnIndex)
            if CString
namePtr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
== CString
forall a. Ptr a
nullPtr
                then SQLError -> IO Text
forall e a. Exception e => e -> IO a
throwIO (Statement -> Int -> SQLError
columnNameUnavailableError Statement
stmt Int
columnIndex)
                else do
                    Text
name <- 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
namePtr
                    Ptr () -> IO ()
c_duckdb_free (CString -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr CString
namePtr)
                    Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
name

-- | Execute a prepared statement and return the number of affected rows.
--   Resets any active result stream before running and raises a 'SqlError'
--   if DuckDB reports a failure.
executeStatement :: Statement -> IO Int
executeStatement :: Statement -> IO Int
executeStatement Statement
stmt =
    Statement -> (DuckDBPreparedStatement -> IO Int) -> IO Int
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
        Statement -> IO ()
resetStatementStream Statement
stmt
        (Ptr DuckDBResult -> IO Int) -> IO Int
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBResult
resPtr -> do
            DuckDBState
rc <- DuckDBPreparedStatement -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_execute_prepared DuckDBPreparedStatement
handle Ptr DuckDBResult
resPtr
            if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                then do
                    Int
changed <- Ptr DuckDBResult -> IO Int
resultRowsChanged Ptr DuckDBResult
resPtr
                    Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                    Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
changed
                else do
                    (Text
errMsg, Maybe DuckDBErrorType
_) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resPtr
                    Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                    SQLError -> IO Int
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO Int) -> SQLError -> IO Int
forall a b. (a -> b) -> a -> b
$ Query -> Text -> SQLError
mkPrepareError (Statement -> Query
statementQuery Statement
stmt) Text
errMsg

-- | Execute a query with positional parameters and return the affected row count.
execute :: (ToRow q) => Connection -> Query -> q -> IO Int
execute :: forall q. ToRow q => Connection -> Query -> q -> IO Int
execute Connection
conn Query
queryText q
params =
    Connection -> Query -> (Statement -> IO Int) -> IO Int
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> [FieldBinding] -> IO ()
bind Statement
stmt (q -> [FieldBinding]
forall a. ToRow a => a -> [FieldBinding]
toRow q
params)
        Statement -> IO Int
executeStatement Statement
stmt

-- | Execute the same query multiple times with different parameter sets.
executeMany :: (ToRow q) => Connection -> Query -> [q] -> IO Int
executeMany :: forall q. ToRow q => Connection -> Query -> [q] -> IO Int
executeMany Connection
conn Query
queryText [q]
rows =
    Connection -> Query -> (Statement -> IO Int) -> IO Int
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> IO [Int] -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (q -> IO Int) -> [q] -> IO [Int]
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 (\q
row -> Statement -> [FieldBinding] -> IO ()
bind Statement
stmt (q -> [FieldBinding]
forall a. ToRow a => a -> [FieldBinding]
toRow q
row) IO () -> IO Int -> IO Int
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Statement -> IO Int
executeStatement Statement
stmt) [q]
rows

-- | Execute an ad-hoc query without parameters and return the affected row count.
execute_ :: Connection -> Query -> IO Int
execute_ :: Connection -> Query -> IO Int
execute_ Connection
conn Query
queryText =
    Connection -> (DuckDBConnection -> IO Int) -> IO Int
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr ->
        Query -> (CString -> IO Int) -> IO Int
forall a. Query -> (CString -> IO a) -> IO a
withQueryCString Query
queryText \CString
sql ->
            (Ptr DuckDBResult -> IO Int) -> IO Int
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 do
                        Int
changed <- Ptr DuckDBResult -> IO Int
resultRowsChanged Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
changed
                    else do
                        (Text
errMsg, Maybe DuckDBErrorType
errType) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        SQLError -> IO Int
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO Int) -> SQLError -> IO Int
forall a b. (a -> b) -> a -> b
$ Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError Query
queryText Text
errMsg Maybe DuckDBErrorType
errType

-- | Execute a query that uses named parameters.
executeNamed :: Connection -> Query -> [NamedParam] -> IO Int
executeNamed :: Connection -> Query -> [NamedParam] -> IO Int
executeNamed Connection
conn Query
queryText [NamedParam]
params =
    Connection -> Query -> (Statement -> IO Int) -> IO Int
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> [NamedParam] -> IO ()
bindNamed Statement
stmt [NamedParam]
params
        Statement -> IO Int
executeStatement Statement
stmt

-- | Run a parameterised query and decode every resulting row eagerly.
query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO [r]
query :: forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query = RowParser r -> Connection -> Query -> q -> IO [r]
forall q r.
ToRow q =>
RowParser r -> Connection -> Query -> q -> IO [r]
queryWith RowParser r
forall a. FromRow a => RowParser a
fromRow

-- | Run a parameterised query with a custom row parser.
queryWith :: (ToRow q) => RowParser r -> Connection -> Query -> q -> IO [r]
queryWith :: forall q r.
ToRow q =>
RowParser r -> Connection -> Query -> q -> IO [r]
queryWith RowParser r
parser Connection
conn Query
queryText q
params =
    Connection -> Query -> (Statement -> IO [r]) -> IO [r]
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> [FieldBinding] -> IO ()
bind Statement
stmt (q -> [FieldBinding]
forall a. ToRow a => a -> [FieldBinding]
toRow q
params)
        Statement -> (DuckDBPreparedStatement -> IO [r]) -> IO [r]
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle ->
            (Ptr DuckDBResult -> IO [r]) -> IO [r]
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBResult
resPtr -> do
                DuckDBState
rc <- DuckDBPreparedStatement -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_execute_prepared DuckDBPreparedStatement
handle Ptr DuckDBResult
resPtr
                if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                    then do
                        [[Field]]
rows <- Ptr DuckDBResult -> IO [[Field]]
collectRows Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        RowParser r -> Query -> [[Field]] -> IO [r]
forall r. RowParser r -> Query -> [[Field]] -> IO [r]
convertRowsWith RowParser r
parser Query
queryText [[Field]]
rows
                    else do
                        (Text
errMsg, Maybe DuckDBErrorType
errType) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        SQLError -> IO [r]
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO [r]) -> SQLError -> IO [r]
forall a b. (a -> b) -> a -> b
$ Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError Query
queryText Text
errMsg Maybe DuckDBErrorType
errType

-- | Run a query that uses named parameters and decode all rows eagerly.
queryNamed :: (FromRow r) => Connection -> Query -> [NamedParam] -> IO [r]
queryNamed :: forall r.
FromRow r =>
Connection -> Query -> [NamedParam] -> IO [r]
queryNamed Connection
conn Query
queryText [NamedParam]
params =
    Connection -> Query -> (Statement -> IO [r]) -> IO [r]
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> [NamedParam] -> IO ()
bindNamed Statement
stmt [NamedParam]
params
        Statement -> (DuckDBPreparedStatement -> IO [r]) -> IO [r]
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle ->
            (Ptr DuckDBResult -> IO [r]) -> IO [r]
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBResult
resPtr -> do
                DuckDBState
rc <- DuckDBPreparedStatement -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_execute_prepared DuckDBPreparedStatement
handle Ptr DuckDBResult
resPtr
                if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                    then do
                        [[Field]]
rows <- Ptr DuckDBResult -> IO [[Field]]
collectRows Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        Query -> [[Field]] -> IO [r]
forall r. FromRow r => Query -> [[Field]] -> IO [r]
convertRows Query
queryText [[Field]]
rows
                    else do
                        (Text
errMsg, Maybe DuckDBErrorType
errType) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        SQLError -> IO [r]
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO [r]) -> SQLError -> IO [r]
forall a b. (a -> b) -> a -> b
$ Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError Query
queryText Text
errMsg Maybe DuckDBErrorType
errType

-- | Run a query without supplying parameters and decode all rows eagerly.
query_ :: (FromRow r) => Connection -> Query -> IO [r]
query_ :: forall r. FromRow r => Connection -> Query -> IO [r]
query_ = RowParser r -> Connection -> Query -> IO [r]
forall r. RowParser r -> Connection -> Query -> IO [r]
queryWith_ RowParser r
forall a. FromRow a => RowParser a
fromRow

-- | Run a query without parameters using a custom row parser.
queryWith_ :: RowParser r -> Connection -> Query -> IO [r]
queryWith_ :: forall r. RowParser r -> Connection -> Query -> IO [r]
queryWith_ RowParser r
parser Connection
conn Query
queryText =
    Connection -> (DuckDBConnection -> IO [r]) -> IO [r]
forall a. Connection -> (DuckDBConnection -> IO a) -> IO a
withConnectionHandle Connection
conn \DuckDBConnection
connPtr ->
        Query -> (CString -> IO [r]) -> IO [r]
forall a. Query -> (CString -> IO a) -> IO a
withQueryCString Query
queryText \CString
sql ->
            (Ptr DuckDBResult -> IO [r]) -> IO [r]
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 do
                        [[Field]]
rows <- Ptr DuckDBResult -> IO [[Field]]
collectRows Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        RowParser r -> Query -> [[Field]] -> IO [r]
forall r. RowParser r -> Query -> [[Field]] -> IO [r]
convertRowsWith RowParser r
parser Query
queryText [[Field]]
rows
                    else do
                        (Text
errMsg, Maybe DuckDBErrorType
errType) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resPtr
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resPtr
                        SQLError -> IO [r]
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO [r]) -> SQLError -> IO [r]
forall a b. (a -> b) -> a -> b
$ Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError Query
queryText Text
errMsg Maybe DuckDBErrorType
errType

-- Streaming folds -----------------------------------------------------------

-- | Stream a parameterised query through an accumulator without loading all rows.
--   Bind the supplied parameters, start a streaming result, and apply the step
--   function row by row to produce a final accumulator value.
fold :: (FromRow row, ToRow params) => Connection -> Query -> params -> a -> (a -> row -> IO a) -> IO a
fold :: forall row params a.
(FromRow row, ToRow params) =>
Connection -> Query -> params -> a -> (a -> row -> IO a) -> IO a
fold Connection
conn Query
queryText params
params a
initial a -> row -> IO a
step =
    Connection -> Query -> (Statement -> IO a) -> IO a
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> IO ()
resetStatementStream Statement
stmt
        Statement -> [FieldBinding] -> IO ()
bind Statement
stmt (params -> [FieldBinding]
forall a. ToRow a => a -> [FieldBinding]
toRow params
params)
        RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
forall row a.
RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
foldStatementWith RowParser row
forall a. FromRow a => RowParser a
fromRow Statement
stmt a
initial a -> row -> IO a
step

-- | Stream a parameterless query through an accumulator without loading all rows.
fold_ :: (FromRow row) => Connection -> Query -> a -> (a -> row -> IO a) -> IO a
fold_ :: forall row a.
FromRow row =>
Connection -> Query -> a -> (a -> row -> IO a) -> IO a
fold_ Connection
conn Query
queryText a
initial a -> row -> IO a
step =
    Connection -> Query -> (Statement -> IO a) -> IO a
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> IO ()
resetStatementStream Statement
stmt
        RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
forall row a.
RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
foldStatementWith RowParser row
forall a. FromRow a => RowParser a
fromRow Statement
stmt a
initial a -> row -> IO a
step

-- | Stream a query that uses named parameters through an accumulator.
foldNamed :: (FromRow row) => Connection -> Query -> [NamedParam] -> a -> (a -> row -> IO a) -> IO a
foldNamed :: forall row a.
FromRow row =>
Connection
-> Query -> [NamedParam] -> a -> (a -> row -> IO a) -> IO a
foldNamed Connection
conn Query
queryText [NamedParam]
params a
initial a -> row -> IO a
step =
    Connection -> Query -> (Statement -> IO a) -> IO a
forall a. Connection -> Query -> (Statement -> IO a) -> IO a
withStatement Connection
conn Query
queryText \Statement
stmt -> do
        Statement -> IO ()
resetStatementStream Statement
stmt
        Statement -> [NamedParam] -> IO ()
bindNamed Statement
stmt [NamedParam]
params
        RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
forall row a.
RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
foldStatementWith RowParser row
forall a. FromRow a => RowParser a
fromRow Statement
stmt a
initial a -> row -> IO a
step

foldStatementWith :: RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
foldStatementWith :: forall row a.
RowParser row -> Statement -> a -> (a -> row -> IO a) -> IO a
foldStatementWith RowParser row
parser Statement
stmt a
initial a -> row -> IO a
step =
    let loop :: a -> IO a
loop a
acc = do
            Maybe row
nextVal <- RowParser row -> Statement -> IO (Maybe row)
forall r. RowParser r -> Statement -> IO (Maybe r)
nextRowWith RowParser row
parser Statement
stmt
            case Maybe row
nextVal of
                Maybe row
Nothing -> a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
acc
                Just row
row -> do
                    a
acc' <- a -> row -> IO a
step a
acc row
row
                    a
acc' a -> IO a -> IO a
forall a b. a -> b -> b
`seq` a -> IO a
loop a
acc'
     in a -> IO a
loop a
initial IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`finally` Statement -> IO ()
resetStatementStream Statement
stmt

-- | Fetch the next row from a streaming statement, stopping when no rows remain.
nextRow :: (FromRow r) => Statement -> IO (Maybe r)
nextRow :: forall r. FromRow r => Statement -> IO (Maybe r)
nextRow = RowParser r -> Statement -> IO (Maybe r)
forall r. RowParser r -> Statement -> IO (Maybe r)
nextRowWith RowParser r
forall a. FromRow a => RowParser a
fromRow

-- | Fetch the next row using a custom parser, returning 'Nothing' once exhausted.
nextRowWith :: RowParser r -> Statement -> IO (Maybe r)
nextRowWith :: forall r. RowParser r -> Statement -> IO (Maybe r)
nextRowWith RowParser r
parser stmt :: Statement
stmt@Statement{IORef StatementStreamState
statementStream :: IORef StatementStreamState
statementStream :: Statement -> IORef StatementStreamState
statementStream} =
    ((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask \forall a. IO a -> IO a
restore -> do
        StatementStreamState
state <- IORef StatementStreamState -> IO StatementStreamState
forall a. IORef a -> IO a
readIORef IORef StatementStreamState
statementStream
        case StatementStreamState
state of
            StatementStreamState
StatementStreamIdle -> do
                Maybe StatementStream
newStream <- IO (Maybe StatementStream) -> IO (Maybe StatementStream)
forall a. IO a -> IO a
restore (Statement -> IO (Maybe StatementStream)
startStatementStream Statement
stmt)
                case Maybe StatementStream
newStream of
                    Maybe StatementStream
Nothing -> Maybe r -> IO (Maybe r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe r
forall a. Maybe a
Nothing
                    Just StatementStream
stream -> IO (Maybe r) -> IO (Maybe r)
forall a. IO a -> IO a
restore (IORef StatementStreamState
-> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
forall r.
IORef StatementStreamState
-> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
consumeStream IORef StatementStreamState
statementStream RowParser r
parser Statement
stmt StatementStream
stream)
            StatementStreamActive StatementStream
stream ->
                IO (Maybe r) -> IO (Maybe r)
forall a. IO a -> IO a
restore (IORef StatementStreamState
-> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
forall r.
IORef StatementStreamState
-> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
consumeStream IORef StatementStreamState
statementStream RowParser r
parser Statement
stmt StatementStream
stream)

resetStatementStream :: Statement -> IO ()
resetStatementStream :: Statement -> IO ()
resetStatementStream Statement{IORef StatementStreamState
statementStream :: Statement -> IORef StatementStreamState
statementStream :: IORef StatementStreamState
statementStream} =
    IORef StatementStreamState -> IO ()
cleanupStatementStreamRef IORef StatementStreamState
statementStream

consumeStream :: IORef StatementStreamState -> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
consumeStream :: forall r.
IORef StatementStreamState
-> RowParser r -> Statement -> StatementStream -> IO (Maybe r)
consumeStream IORef StatementStreamState
streamRef RowParser r
parser Statement
stmt StatementStream
stream = do
    Either SomeException (Maybe [Field], StatementStream)
result <-
        ( IO (Maybe [Field], StatementStream)
-> IO (Either SomeException (Maybe [Field], StatementStream))
forall e a. Exception e => IO a -> IO (Either e a)
try (Query -> StatementStream -> IO (Maybe [Field], StatementStream)
streamNextRow (Statement -> Query
statementQuery Statement
stmt) StatementStream
stream)
            :: IO (Either SomeException (Maybe [Field], StatementStream))
        )
    case Either SomeException (Maybe [Field], StatementStream)
result of
        Left SomeException
err -> do
            StatementStream -> IO ()
finalizeStream StatementStream
stream
            IORef StatementStreamState -> StatementStreamState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef StatementStreamState
streamRef StatementStreamState
StatementStreamIdle
            SomeException -> IO (Maybe r)
forall e a. Exception e => e -> IO a
throwIO SomeException
err
        Right (Maybe [Field]
maybeFields, StatementStream
updatedStream) ->
            case Maybe [Field]
maybeFields of
                Maybe [Field]
Nothing -> do
                    StatementStream -> IO ()
finalizeStream StatementStream
updatedStream
                    IORef StatementStreamState -> StatementStreamState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef StatementStreamState
streamRef StatementStreamState
StatementStreamIdle
                    Maybe r -> IO (Maybe r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe r
forall a. Maybe a
Nothing
                Just [Field]
fields ->
                    case RowParser r -> [Field] -> Ok r
forall a. RowParser a -> [Field] -> Ok a
parseRow RowParser r
parser [Field]
fields of
                        Errors [SomeException]
rowErr -> do
                            StatementStream -> IO ()
finalizeStream StatementStream
updatedStream
                            IORef StatementStreamState -> StatementStreamState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef StatementStreamState
streamRef StatementStreamState
StatementStreamIdle
                            SQLError -> IO (Maybe r)
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO (Maybe r)) -> SQLError -> IO (Maybe r)
forall a b. (a -> b) -> a -> b
$ Query -> [SomeException] -> SQLError
rowErrorsToSqlError (Statement -> Query
statementQuery Statement
stmt) [SomeException]
rowErr
                        Ok r
value -> do
                            IORef StatementStreamState -> StatementStreamState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef StatementStreamState
streamRef (StatementStream -> StatementStreamState
StatementStreamActive StatementStream
updatedStream)
                            Maybe r -> IO (Maybe r)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (r -> Maybe r
forall a. a -> Maybe a
Just r
value)

startStatementStream :: Statement -> IO (Maybe StatementStream)
startStatementStream :: Statement -> IO (Maybe StatementStream)
startStatementStream Statement
stmt =
    Statement
-> (DuckDBPreparedStatement -> IO (Maybe StatementStream))
-> IO (Maybe StatementStream)
forall a. Statement -> (DuckDBPreparedStatement -> IO a) -> IO a
withStatementHandle Statement
stmt \DuckDBPreparedStatement
handle -> do
        [StatementStreamColumn]
columns <- DuckDBPreparedStatement -> IO [StatementStreamColumn]
collectStreamColumns DuckDBPreparedStatement
handle
        Ptr DuckDBResult
resultPtr <- IO (Ptr DuckDBResult)
forall a. Storable a => IO (Ptr a)
malloc
        DuckDBState
rc <- DuckDBPreparedStatement -> Ptr DuckDBResult -> IO DuckDBState
c_duckdb_execute_prepared_streaming DuckDBPreparedStatement
handle Ptr DuckDBResult
resultPtr
        if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
/= DuckDBState
DuckDBSuccess
            then do
                (Text
errMsg, Maybe DuckDBErrorType
errType) <- Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resultPtr
                Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resultPtr
                Ptr DuckDBResult -> IO ()
forall a. Ptr a -> IO ()
free Ptr DuckDBResult
resultPtr
                SQLError -> IO (Maybe StatementStream)
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO (Maybe StatementStream))
-> SQLError -> IO (Maybe StatementStream)
forall a b. (a -> b) -> a -> b
$ Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError (Statement -> Query
statementQuery Statement
stmt) Text
errMsg Maybe DuckDBErrorType
errType
            else do
                DuckDBResultType
resultType <- Ptr DuckDBResult -> IO DuckDBResultType
c_duckdb_result_return_type Ptr DuckDBResult
resultPtr
                if DuckDBResultType
resultType DuckDBResultType -> DuckDBResultType -> Bool
forall a. Eq a => a -> a -> Bool
/= DuckDBResultType
DuckDBResultTypeQueryResult
                    then do
                        Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
resultPtr
                        Ptr DuckDBResult -> IO ()
forall a. Ptr a -> IO ()
free Ptr DuckDBResult
resultPtr
                        Maybe StatementStream -> IO (Maybe StatementStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StatementStream
forall a. Maybe a
Nothing
                    else Maybe StatementStream -> IO (Maybe StatementStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StatementStream -> Maybe StatementStream
forall a. a -> Maybe a
Just (Ptr DuckDBResult
-> [StatementStreamColumn]
-> Maybe StatementStreamChunk
-> StatementStream
StatementStream Ptr DuckDBResult
resultPtr [StatementStreamColumn]
columns Maybe StatementStreamChunk
forall a. Maybe a
Nothing))

collectStreamColumns :: DuckDBPreparedStatement -> IO [StatementStreamColumn]
collectStreamColumns :: DuckDBPreparedStatement -> IO [StatementStreamColumn]
collectStreamColumns DuckDBPreparedStatement
handle = do
    DuckDBIdx
rawCount <- DuckDBPreparedStatement -> IO DuckDBIdx
c_duckdb_prepared_statement_column_count DuckDBPreparedStatement
handle
    let cc :: Int
cc = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
rawCount :: Int
    [Int]
-> (Int -> IO StatementStreamColumn) -> IO [StatementStreamColumn]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] \Int
idx -> do
        CString
namePtr <- DuckDBPreparedStatement -> DuckDBIdx -> IO CString
c_duckdb_prepared_statement_column_name DuckDBPreparedStatement
handle (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx)
        Text
name <-
            if CString
namePtr 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
"column" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
idx))
                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
namePtr
        DuckDBType
dtype <- DuckDBPreparedStatement -> DuckDBIdx -> IO DuckDBType
c_duckdb_prepared_statement_column_type DuckDBPreparedStatement
handle (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx)
        StatementStreamColumn -> IO StatementStreamColumn
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            StatementStreamColumn
                { statementStreamColumnIndex :: Int
statementStreamColumnIndex = Int
idx
                , statementStreamColumnName :: Text
statementStreamColumnName = Text
name
                , statementStreamColumnType :: DuckDBType
statementStreamColumnType = DuckDBType
dtype
                }

streamNextRow :: Query -> StatementStream -> IO (Maybe [Field], StatementStream)
streamNextRow :: Query -> StatementStream -> IO (Maybe [Field], StatementStream)
streamNextRow Query
queryText stream :: StatementStream
stream@StatementStream{statementStreamChunk :: StatementStream -> Maybe StatementStreamChunk
statementStreamChunk = Maybe StatementStreamChunk
Nothing} = do
    StatementStream
refreshed <- StatementStream -> IO StatementStream
fetchChunk StatementStream
stream
    case StatementStream -> Maybe StatementStreamChunk
statementStreamChunk StatementStream
refreshed of
        Maybe StatementStreamChunk
Nothing -> (Maybe [Field], StatementStream)
-> IO (Maybe [Field], StatementStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [Field]
forall a. Maybe a
Nothing, StatementStream
refreshed)
        Just StatementStreamChunk
chunk -> Query
-> StatementStream
-> StatementStreamChunk
-> IO (Maybe [Field], StatementStream)
emitRow Query
queryText StatementStream
refreshed StatementStreamChunk
chunk
streamNextRow Query
queryText stream :: StatementStream
stream@StatementStream{statementStreamChunk :: StatementStream -> Maybe StatementStreamChunk
statementStreamChunk = Just StatementStreamChunk
chunk} =
    Query
-> StatementStream
-> StatementStreamChunk
-> IO (Maybe [Field], StatementStream)
emitRow Query
queryText StatementStream
stream StatementStreamChunk
chunk

fetchChunk :: StatementStream -> IO StatementStream
fetchChunk :: StatementStream -> IO StatementStream
fetchChunk stream :: StatementStream
stream@StatementStream{Ptr DuckDBResult
statementStreamResult :: Ptr DuckDBResult
statementStreamResult :: StatementStream -> Ptr DuckDBResult
statementStreamResult} = do
    DuckDBDataChunk
chunk <- Ptr DuckDBResult -> IO DuckDBDataChunk
c_duckdb_stream_fetch_chunk Ptr DuckDBResult
statementStreamResult
    if DuckDBDataChunk
chunk DuckDBDataChunk -> DuckDBDataChunk -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBDataChunk
forall a. Ptr a
nullPtr
        then StatementStream -> IO StatementStream
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StatementStream
stream
        else do
            DuckDBIdx
rawSize <- DuckDBDataChunk -> IO DuckDBIdx
c_duckdb_data_chunk_get_size DuckDBDataChunk
chunk
            let rowCount :: Int
rowCount = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
rawSize :: Int
            if Int
rowCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
                then do
                    DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk
                    StatementStream -> IO StatementStream
fetchChunk StatementStream
stream
                else do
                    [StatementStreamChunkVector]
vectors <- DuckDBDataChunk
-> [StatementStreamColumn] -> IO [StatementStreamChunkVector]
prepareChunkVectors DuckDBDataChunk
chunk (StatementStream -> [StatementStreamColumn]
statementStreamColumns StatementStream
stream)
                    let chunkState :: StatementStreamChunk
chunkState =
                            StatementStreamChunk
                                { statementStreamChunkPtr :: DuckDBDataChunk
statementStreamChunkPtr = DuckDBDataChunk
chunk
                                , statementStreamChunkSize :: Int
statementStreamChunkSize = Int
rowCount
                                , statementStreamChunkIndex :: Int
statementStreamChunkIndex = Int
0
                                , statementStreamChunkVectors :: [StatementStreamChunkVector]
statementStreamChunkVectors = [StatementStreamChunkVector]
vectors
                                }
                    StatementStream -> IO StatementStream
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StatementStream
stream{statementStreamChunk = Just chunkState}

prepareChunkVectors :: DuckDBDataChunk -> [StatementStreamColumn] -> IO [StatementStreamChunkVector]
prepareChunkVectors :: DuckDBDataChunk
-> [StatementStreamColumn] -> IO [StatementStreamChunkVector]
prepareChunkVectors DuckDBDataChunk
chunk [StatementStreamColumn]
columns =
    [StatementStreamColumn]
-> (StatementStreamColumn -> IO StatementStreamChunkVector)
-> IO [StatementStreamChunkVector]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [StatementStreamColumn]
columns \StatementStreamColumn{Int
statementStreamColumnIndex :: StatementStreamColumn -> Int
statementStreamColumnIndex :: Int
statementStreamColumnIndex} -> do
        DuckDBVector
vector <- DuckDBDataChunk -> DuckDBIdx -> IO DuckDBVector
c_duckdb_data_chunk_get_vector DuckDBDataChunk
chunk (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
statementStreamColumnIndex)
        Ptr ()
dataPtr <- DuckDBVector -> IO (Ptr ())
c_duckdb_vector_get_data DuckDBVector
vector
        Ptr DuckDBIdx
validity <- DuckDBVector -> IO (Ptr DuckDBIdx)
c_duckdb_vector_get_validity DuckDBVector
vector
        StatementStreamChunkVector -> IO StatementStreamChunkVector
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            StatementStreamChunkVector
                { statementStreamChunkVectorHandle :: DuckDBVector
statementStreamChunkVectorHandle = DuckDBVector
vector
                , statementStreamChunkVectorData :: Ptr ()
statementStreamChunkVectorData = Ptr ()
dataPtr
                , statementStreamChunkVectorValidity :: Ptr DuckDBIdx
statementStreamChunkVectorValidity = Ptr DuckDBIdx
validity
                }

emitRow :: Query -> StatementStream -> StatementStreamChunk -> IO (Maybe [Field], StatementStream)
emitRow :: Query
-> StatementStream
-> StatementStreamChunk
-> IO (Maybe [Field], StatementStream)
emitRow Query
queryText StatementStream
stream chunk :: StatementStreamChunk
chunk@StatementStreamChunk{Int
statementStreamChunkIndex :: StatementStreamChunk -> Int
statementStreamChunkIndex :: Int
statementStreamChunkIndex, Int
statementStreamChunkSize :: StatementStreamChunk -> Int
statementStreamChunkSize :: Int
statementStreamChunkSize} = do
    [Field]
fields <-
        Query
-> [StatementStreamColumn]
-> [StatementStreamChunkVector]
-> Int
-> IO [Field]
buildRow
            Query
queryText
            (StatementStream -> [StatementStreamColumn]
statementStreamColumns StatementStream
stream)
            (StatementStreamChunk -> [StatementStreamChunkVector]
statementStreamChunkVectors StatementStreamChunk
chunk)
            Int
statementStreamChunkIndex
    let nextIndex :: Int
nextIndex = Int
statementStreamChunkIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    if Int
nextIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
statementStreamChunkSize
        then
            let updatedChunk :: StatementStreamChunk
updatedChunk = StatementStreamChunk
chunk{statementStreamChunkIndex = nextIndex}
             in (Maybe [Field], StatementStream)
-> IO (Maybe [Field], StatementStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Field] -> Maybe [Field]
forall a. a -> Maybe a
Just [Field]
fields, StatementStream
stream{statementStreamChunk = Just updatedChunk})
        else do
            DuckDBDataChunk -> IO ()
destroyDataChunk (StatementStreamChunk -> DuckDBDataChunk
statementStreamChunkPtr StatementStreamChunk
chunk)
            (Maybe [Field], StatementStream)
-> IO (Maybe [Field], StatementStream)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Field] -> Maybe [Field]
forall a. a -> Maybe a
Just [Field]
fields, StatementStream
stream{statementStreamChunk = Nothing})

buildRow :: Query -> [StatementStreamColumn] -> [StatementStreamChunkVector] -> Int -> IO [Field]
buildRow :: Query
-> [StatementStreamColumn]
-> [StatementStreamChunkVector]
-> Int
-> IO [Field]
buildRow Query
queryText [StatementStreamColumn]
columns [StatementStreamChunkVector]
vectors Int
rowIdx =
    (StatementStreamColumn -> StatementStreamChunkVector -> IO Field)
-> [StatementStreamColumn]
-> [StatementStreamChunkVector]
-> IO [Field]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (Query
-> Int
-> StatementStreamColumn
-> StatementStreamChunkVector
-> IO Field
buildField Query
queryText Int
rowIdx) [StatementStreamColumn]
columns [StatementStreamChunkVector]
vectors

buildField :: Query -> Int -> StatementStreamColumn -> StatementStreamChunkVector -> IO Field
buildField :: Query
-> Int
-> StatementStreamColumn
-> StatementStreamChunkVector
-> IO Field
buildField Query
queryText Int
rowIdx StatementStreamColumn
column StatementStreamChunkVector{Ptr ()
statementStreamChunkVectorData :: StatementStreamChunkVector -> Ptr ()
statementStreamChunkVectorData :: Ptr ()
statementStreamChunkVectorData, Ptr DuckDBIdx
statementStreamChunkVectorValidity :: StatementStreamChunkVector -> Ptr DuckDBIdx
statementStreamChunkVectorValidity :: Ptr DuckDBIdx
statementStreamChunkVectorValidity} = do
    let duckIdx :: DuckDBIdx
duckIdx = Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rowIdx
        dtype :: DuckDBType
dtype = StatementStreamColumn -> DuckDBType
statementStreamColumnType StatementStreamColumn
column
    Bool
valid <- Ptr DuckDBIdx -> DuckDBIdx -> IO Bool
chunkIsRowValid Ptr DuckDBIdx
statementStreamChunkVectorValidity DuckDBIdx
duckIdx
    FieldValue
value <-
        if Bool -> Bool
not Bool
valid
            then FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FieldValue
FieldNull
            else case DuckDBType
dtype of
                DuckDBType
DuckDBTypeBoolean -> do
                    Word8
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Word8) Int
rowIdx
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> FieldValue
FieldBool (Word8
raw Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0))
                DuckDBType
DuckDBTypeTinyInt -> do
                    Int8
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Int8) Int
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
raw))
                DuckDBType
DuckDBTypeSmallInt -> do
                    Int16
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Int16) Int
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
raw))
                DuckDBType
DuckDBTypeInteger -> do
                    Int32
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Int32) Int
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
raw))
                DuckDBType
DuckDBTypeBigInt -> do
                    Int64
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Int64) Int
rowIdx
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> FieldValue
FieldInt64 Int64
raw)
                DuckDBType
DuckDBTypeUTinyInt -> do
                    Word8
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Word8) Int
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
raw))
                DuckDBType
DuckDBTypeUSmallInt -> do
                    Word16
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Word16) Int
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
raw))
                DuckDBType
DuckDBTypeUInteger -> do
                    Word32
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Word32) Int
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
raw))
                DuckDBType
DuckDBTypeUBigInt -> do
                    DuckDBIdx
raw <- Ptr DuckDBIdx -> Int -> IO DuckDBIdx
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBIdx
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
statementStreamChunkVectorData :: Ptr Word64) Int
rowIdx
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DuckDBIdx -> FieldValue
FieldWord64 (DuckDBIdx -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
raw))
                DuckDBType
DuckDBTypeBlob -> do
                    ByteString
blob <- Ptr () -> DuckDBIdx -> IO ByteString
chunkDecodeBlob Ptr ()
statementStreamChunkVectorData DuckDBIdx
duckIdx
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> FieldValue
FieldBlob ByteString
blob)
                DuckDBType
DuckDBTypeDate -> do
                    Int32
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Int32) Int
rowIdx
                    Day
day <- DuckDBDate -> IO Day
decodeDuckDBDate (Int32 -> DuckDBDate
DuckDBDate Int32
raw)
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Day -> FieldValue
FieldDate Day
day)
                DuckDBType
DuckDBTypeTime -> do
                    DuckDBTime
raw <- Ptr DuckDBTime -> Int -> IO DuckDBTime
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBTime
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
statementStreamChunkVectorData :: Ptr DuckDBTime) Int
rowIdx
                    TimeOfDay
tod <- DuckDBTime -> IO TimeOfDay
decodeDuckDBTime DuckDBTime
raw
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TimeOfDay -> FieldValue
FieldTime TimeOfDay
tod)
                DuckDBType
DuckDBTypeTimestamp -> do
                    DuckDBTimestamp
raw <- Ptr DuckDBTimestamp -> Int -> IO DuckDBTimestamp
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBTimestamp
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
statementStreamChunkVectorData :: Ptr DuckDBTimestamp) Int
rowIdx
                    LocalTime
ts <- DuckDBTimestamp -> IO LocalTime
decodeDuckDBTimestamp DuckDBTimestamp
raw
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LocalTime -> FieldValue
FieldTimestamp LocalTime
ts)
                DuckDBType
DuckDBTypeFloat -> do
                    Float
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Float) Int
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
raw))
                DuckDBType
DuckDBTypeDouble -> do
                    Double
raw <- 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 ()
statementStreamChunkVectorData :: Ptr Double) Int
rowIdx
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> FieldValue
FieldDouble Double
raw)
                DuckDBType
DuckDBTypeVarchar -> Text -> FieldValue
FieldText (Text -> FieldValue) -> IO Text -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr () -> DuckDBIdx -> IO Text
chunkDecodeText Ptr ()
statementStreamChunkVectorData DuckDBIdx
duckIdx
                DuckDBType
DuckDBTypeHugeInt -> String -> IO FieldValue
forall a. HasCallStack => String -> a
error String
"duckdb-simple: HUGEINT is not supported"
                DuckDBType
DuckDBTypeUHugeInt -> String -> IO FieldValue
forall a. HasCallStack => String -> a
error String
"duckdb-simple: HUGEINT is not supported"
                DuckDBType
_ ->
                    SQLError -> IO FieldValue
forall e a. Exception e => e -> IO a
throwIO (Query -> StatementStreamColumn -> SQLError
streamingUnsupportedTypeError Query
queryText StatementStreamColumn
column)
    Field -> IO Field
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        Field
            { fieldName :: Text
fieldName = StatementStreamColumn -> Text
statementStreamColumnName StatementStreamColumn
column
            , fieldIndex :: Int
fieldIndex = StatementStreamColumn -> Int
statementStreamColumnIndex StatementStreamColumn
column
            , fieldValue :: FieldValue
fieldValue = FieldValue
value
            }

chunkIsRowValid :: Ptr Word64 -> DuckDBIdx -> IO Bool
chunkIsRowValid :: Ptr DuckDBIdx -> DuckDBIdx -> IO Bool
chunkIsRowValid Ptr DuckDBIdx
validity DuckDBIdx
rowIdx
    | Ptr DuckDBIdx
validity Ptr DuckDBIdx -> Ptr DuckDBIdx -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr DuckDBIdx
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 DuckDBIdx -> DuckDBIdx -> IO CBool
c_duckdb_validity_row_is_valid Ptr DuckDBIdx
validity DuckDBIdx
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)

chunkDecodeText :: Ptr () -> DuckDBIdx -> IO Text
chunkDecodeText :: Ptr () -> DuckDBIdx -> IO Text
chunkDecodeText Ptr ()
dataPtr DuckDBIdx
rowIdx = do
    let base :: Ptr Word8
base = Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8
        offset :: Int
offset = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
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
base 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
TextEncoding.decodeUtf8 ByteString
bytes)

chunkDecodeBlob :: Ptr () -> DuckDBIdx -> IO BS.ByteString
chunkDecodeBlob :: Ptr () -> DuckDBIdx -> IO ByteString
chunkDecodeBlob Ptr ()
dataPtr DuckDBIdx
rowIdx = do
    let base :: Ptr Word8
base = Ptr () -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word8
        offset :: Int
offset = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
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
base 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 ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
        else do
            CString
ptr <- Ptr DuckDBStringT -> IO CString
c_duckdb_string_t_data Ptr DuckDBStringT
stringPtr
            CStringLen -> IO ByteString
BS.packCStringLen (CString
ptr, Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)

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

cleanupStatementStreamRef :: IORef StatementStreamState -> IO ()
cleanupStatementStreamRef :: IORef StatementStreamState -> IO ()
cleanupStatementStreamRef IORef StatementStreamState
ref = do
    StatementStreamState
state <- IORef StatementStreamState
-> (StatementStreamState
    -> (StatementStreamState, StatementStreamState))
-> IO StatementStreamState
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef StatementStreamState
ref (StatementStreamState
StatementStreamIdle,)
    StatementStreamState -> IO ()
finalizeStreamState StatementStreamState
state

finalizeStreamState :: StatementStreamState -> IO ()
finalizeStreamState :: StatementStreamState -> IO ()
finalizeStreamState = \case
    StatementStreamState
StatementStreamIdle -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    StatementStreamActive StatementStream
stream -> StatementStream -> IO ()
finalizeStream StatementStream
stream

finalizeStream :: StatementStream -> IO ()
finalizeStream :: StatementStream -> IO ()
finalizeStream StatementStream{Ptr DuckDBResult
statementStreamResult :: StatementStream -> Ptr DuckDBResult
statementStreamResult :: Ptr DuckDBResult
statementStreamResult, Maybe StatementStreamChunk
statementStreamChunk :: StatementStream -> Maybe StatementStreamChunk
statementStreamChunk :: Maybe StatementStreamChunk
statementStreamChunk} = do
    IO ()
-> (StatementStreamChunk -> IO ())
-> Maybe StatementStreamChunk
-> IO ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) StatementStreamChunk -> IO ()
finalizeChunk Maybe StatementStreamChunk
statementStreamChunk
    Ptr DuckDBResult -> IO ()
c_duckdb_destroy_result Ptr DuckDBResult
statementStreamResult
    Ptr DuckDBResult -> IO ()
forall a. Ptr a -> IO ()
free Ptr DuckDBResult
statementStreamResult

finalizeChunk :: StatementStreamChunk -> IO ()
finalizeChunk :: StatementStreamChunk -> IO ()
finalizeChunk StatementStreamChunk{DuckDBDataChunk
statementStreamChunkPtr :: StatementStreamChunk -> DuckDBDataChunk
statementStreamChunkPtr :: DuckDBDataChunk
statementStreamChunkPtr} =
    DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
statementStreamChunkPtr

destroyDataChunk :: DuckDBDataChunk -> IO ()
destroyDataChunk :: DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk =
    (Ptr DuckDBDataChunk -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBDataChunk
ptr -> do
        Ptr DuckDBDataChunk -> DuckDBDataChunk -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBDataChunk
ptr DuckDBDataChunk
chunk
        Ptr DuckDBDataChunk -> IO ()
c_duckdb_destroy_data_chunk Ptr DuckDBDataChunk
ptr

streamingUnsupportedTypeError :: Query -> StatementStreamColumn -> SQLError
streamingUnsupportedTypeError :: Query -> StatementStreamColumn -> SQLError
streamingUnsupportedTypeError Query
queryText StatementStreamColumn{Text
statementStreamColumnName :: StatementStreamColumn -> Text
statementStreamColumnName :: Text
statementStreamColumnName, DuckDBType
statementStreamColumnType :: StatementStreamColumn -> DuckDBType
statementStreamColumnType :: DuckDBType
statementStreamColumnType} =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage =
            [Text] -> Text
Text.concat
                [ String -> Text
Text.pack String
"duckdb-simple: streaming does not yet support column "
                , Text
statementStreamColumnName
                , String -> Text
Text.pack String
" with DuckDB type "
                , String -> Text
Text.pack (DuckDBType -> String
forall a. Show a => a -> String
show DuckDBType
statementStreamColumnType)
                ]
        , 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
queryText
        }

-- | Run an action inside a transaction.
withTransaction :: Connection -> IO a -> IO a
withTransaction :: forall a. Connection -> IO a -> IO a
withTransaction Connection
conn IO a
action =
    ((forall a. IO a -> IO a) -> IO a) -> IO a
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask \forall a. IO a -> IO a
restore -> do
        IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection -> Query -> IO Int
execute_ Connection
conn Query
begin)
        let rollbackAction :: IO ()
rollbackAction = IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection -> Query -> IO Int
execute_ Connection
conn Query
rollback)
        a
result <- IO a -> IO a
forall a. IO a -> IO a
restore IO a
action IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` IO ()
rollbackAction
        IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Connection -> Query -> IO Int
execute_ Connection
conn Query
commit)
        a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
result

  where begin :: Query
begin = Text -> Query
Query (String -> Text
Text.pack String
"BEGIN TRANSACTION")
        commit :: Query
commit = Text -> Query
Query (String -> Text
Text.pack String
"COMMIT")
        rollback :: Query
rollback = Text -> Query
Query (String -> Text
Text.pack String
"ROLLBACK")


-- Internal helpers -----------------------------------------------------------

createConnection :: DuckDBDatabase -> DuckDBConnection -> IO Connection
createConnection :: DuckDBDatabase -> DuckDBConnection -> IO Connection
createConnection DuckDBDatabase
db DuckDBConnection
conn = do
    IORef ConnectionState
ref <- ConnectionState -> IO (IORef ConnectionState)
forall a. a -> IO (IORef a)
newIORef (DuckDBDatabase -> DuckDBConnection -> ConnectionState
ConnectionOpen DuckDBDatabase
db DuckDBConnection
conn)
    Weak (IORef ConnectionState)
_ <-
        IORef ConnectionState -> IO () -> IO (Weak (IORef ConnectionState))
forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ConnectionState
ref (IO () -> IO (Weak (IORef ConnectionState)))
-> IO () -> IO (Weak (IORef ConnectionState))
forall a b. (a -> b) -> a -> b
$
            IO (IO ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
                IORef ConnectionState
-> (ConnectionState -> (ConnectionState, IO ())) -> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef ConnectionState
ref \case
                    ConnectionState
ConnectionClosed -> (ConnectionState
ConnectionClosed, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
                    openState :: ConnectionState
openState@(ConnectionOpen{}) ->
                        (ConnectionState
ConnectionClosed, ConnectionState -> IO ()
closeHandles ConnectionState
openState)
    Connection -> IO Connection
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Connection{connectionState :: IORef ConnectionState
connectionState = IORef ConnectionState
ref}

createStatement :: Connection -> DuckDBPreparedStatement -> Query -> IO Statement
createStatement :: Connection -> DuckDBPreparedStatement -> Query -> IO Statement
createStatement Connection
parent DuckDBPreparedStatement
handle Query
queryText = do
    IORef StatementState
ref <- StatementState -> IO (IORef StatementState)
forall a. a -> IO (IORef a)
newIORef (DuckDBPreparedStatement -> StatementState
StatementOpen DuckDBPreparedStatement
handle)
    IORef StatementStreamState
streamRef <- StatementStreamState -> IO (IORef StatementStreamState)
forall a. a -> IO (IORef a)
newIORef StatementStreamState
StatementStreamIdle
    Weak (IORef StatementState)
_ <-
        IORef StatementState -> IO () -> IO (Weak (IORef StatementState))
forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef StatementState
ref (IO () -> IO (Weak (IORef StatementState)))
-> IO () -> IO (Weak (IORef StatementState))
forall a b. (a -> b) -> a -> b
$
          do IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
               IORef StatementState
-> (StatementState -> (StatementState, IO ())) -> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef StatementState
ref ((StatementState -> (StatementState, IO ())) -> IO (IO ()))
-> (StatementState -> (StatementState, IO ())) -> IO (IO ())
forall a b. (a -> b) -> a -> b
$ \case
                        StatementState
StatementClosed -> (StatementState
StatementClosed, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
                        StatementOpen{DuckDBPreparedStatement
statementHandle :: StatementState -> DuckDBPreparedStatement
statementHandle :: DuckDBPreparedStatement
statementHandle} ->
                            ( StatementState
StatementClosed
                            , do
                                IORef StatementStreamState -> IO ()
cleanupStatementStreamRef IORef StatementStreamState
streamRef
                                DuckDBPreparedStatement -> IO ()
destroyPrepared DuckDBPreparedStatement
statementHandle
                            )
    Statement -> IO Statement
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        Statement
            { statementState :: IORef StatementState
statementState = IORef StatementState
ref
            , statementConnection :: Connection
statementConnection = Connection
parent
            , statementQuery :: Query
statementQuery = Query
queryText
            , statementStream :: IORef StatementStreamState
statementStream = IORef StatementStreamState
streamRef
            }

openDatabase :: FilePath -> IO DuckDBDatabase
openDatabase :: String -> IO DuckDBDatabase
openDatabase String
path =
    (Ptr DuckDBDatabase -> IO DuckDBDatabase) -> IO DuckDBDatabase
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBDatabase
dbPtr ->
        (Ptr CString -> IO DuckDBDatabase) -> IO DuckDBDatabase
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr CString
errPtr -> do
            Ptr CString -> CString -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr CString
errPtr CString
forall a. Ptr a
nullPtr
            String -> (CString -> IO DuckDBDatabase) -> IO DuckDBDatabase
forall a. String -> (CString -> IO a) -> IO a
withCString String
path \CString
cPath -> do
                DuckDBState
rc <- CString
-> Ptr DuckDBDatabase
-> DuckDBConfig
-> Ptr CString
-> IO DuckDBState
c_duckdb_open_ext CString
cPath Ptr DuckDBDatabase
dbPtr DuckDBConfig
forall a. Ptr a
nullPtr Ptr CString
errPtr
                if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
                    then do
                        DuckDBDatabase
db <- Ptr DuckDBDatabase -> IO DuckDBDatabase
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBDatabase
dbPtr
                        Ptr CString -> IO ()
maybeFreeErr Ptr CString
errPtr
                        DuckDBDatabase -> IO DuckDBDatabase
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DuckDBDatabase
db
                    else do
                        Text
errMsg <- Ptr CString -> IO Text
peekError Ptr CString
errPtr
                        Ptr CString -> IO ()
maybeFreeErr Ptr CString
errPtr
                        SQLError -> IO DuckDBDatabase
forall e a. Exception e => e -> IO a
throwIO (SQLError -> IO DuckDBDatabase) -> SQLError -> IO DuckDBDatabase
forall a b. (a -> b) -> a -> b
$ Text -> SQLError
mkOpenError Text
errMsg

connectDatabase :: DuckDBDatabase -> IO DuckDBConnection
connectDatabase :: DuckDBDatabase -> IO DuckDBConnection
connectDatabase DuckDBDatabase
db =
    (Ptr DuckDBConnection -> IO DuckDBConnection)
-> IO DuckDBConnection
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBConnection
connPtr -> do
        DuckDBState
rc <- DuckDBDatabase -> Ptr DuckDBConnection -> IO DuckDBState
c_duckdb_connect DuckDBDatabase
db Ptr DuckDBConnection
connPtr
        if DuckDBState
rc DuckDBState -> DuckDBState -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBState
DuckDBSuccess
            then Ptr DuckDBConnection -> IO DuckDBConnection
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBConnection
connPtr
            else SQLError -> IO DuckDBConnection
forall e a. Exception e => e -> IO a
throwIO SQLError
mkConnectError

closeHandles :: ConnectionState -> IO ()
closeHandles :: ConnectionState -> IO ()
closeHandles ConnectionState
ConnectionClosed = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
closeHandles ConnectionOpen{DuckDBDatabase
connectionDatabase :: DuckDBDatabase
connectionDatabase :: ConnectionState -> DuckDBDatabase
connectionDatabase, DuckDBConnection
connectionHandle :: DuckDBConnection
connectionHandle :: ConnectionState -> DuckDBConnection
connectionHandle} = do
    DuckDBConnection -> IO ()
closeConnectionHandle DuckDBConnection
connectionHandle
    DuckDBDatabase -> IO ()
closeDatabaseHandle DuckDBDatabase
connectionDatabase

closeConnectionHandle :: DuckDBConnection -> IO ()
closeConnectionHandle :: DuckDBConnection -> IO ()
closeConnectionHandle DuckDBConnection
conn =
    (Ptr DuckDBConnection -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBConnection
ptr -> Ptr DuckDBConnection -> DuckDBConnection -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBConnection
ptr DuckDBConnection
conn IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr DuckDBConnection -> IO ()
c_duckdb_disconnect Ptr DuckDBConnection
ptr

closeDatabaseHandle :: DuckDBDatabase -> IO ()
closeDatabaseHandle :: DuckDBDatabase -> IO ()
closeDatabaseHandle DuckDBDatabase
db =
    (Ptr DuckDBDatabase -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBDatabase
ptr -> Ptr DuckDBDatabase -> DuckDBDatabase -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBDatabase
ptr DuckDBDatabase
db IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr DuckDBDatabase -> IO ()
c_duckdb_close Ptr DuckDBDatabase
ptr

destroyPrepared :: DuckDBPreparedStatement -> IO ()
destroyPrepared :: DuckDBPreparedStatement -> IO ()
destroyPrepared DuckDBPreparedStatement
stmt =
    (Ptr DuckDBPreparedStatement -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBPreparedStatement
ptr -> Ptr DuckDBPreparedStatement -> DuckDBPreparedStatement -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBPreparedStatement
ptr DuckDBPreparedStatement
stmt IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr DuckDBPreparedStatement -> IO ()
c_duckdb_destroy_prepare Ptr DuckDBPreparedStatement
ptr

fetchPrepareError :: DuckDBPreparedStatement -> IO Text
fetchPrepareError :: DuckDBPreparedStatement -> IO Text
fetchPrepareError DuckDBPreparedStatement
stmt = do
    CString
msgPtr <- DuckDBPreparedStatement -> IO CString
c_duckdb_prepare_error DuckDBPreparedStatement
stmt
    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: prepare 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

fetchResultError :: Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError :: Ptr DuckDBResult -> IO (Text, Maybe DuckDBErrorType)
fetchResultError Ptr DuckDBResult
resultPtr = do
    CString
msgPtr <- Ptr DuckDBResult -> IO CString
c_duckdb_result_error Ptr DuckDBResult
resultPtr
    Text
msg <-
        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: query 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
    DuckDBErrorType
errType <- Ptr DuckDBResult -> IO DuckDBErrorType
c_duckdb_result_error_type Ptr DuckDBResult
resultPtr
    let classified :: Maybe DuckDBErrorType
classified =
            if DuckDBErrorType
errType DuckDBErrorType -> DuckDBErrorType -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBErrorType
DuckDBErrorInvalid
                then Maybe DuckDBErrorType
forall a. Maybe a
Nothing
                else DuckDBErrorType -> Maybe DuckDBErrorType
forall a. a -> Maybe a
Just DuckDBErrorType
errType
    (Text, Maybe DuckDBErrorType) -> IO (Text, Maybe DuckDBErrorType)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
msg, Maybe DuckDBErrorType
classified)

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

mkConnectError :: SQLError
mkConnectError :: SQLError
mkConnectError =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage = String -> Text
Text.pack String
"duckdb-simple: failed to create connection handle"
        , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
forall a. Maybe a
Nothing
        , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Maybe Query
forall a. Maybe a
Nothing
        }

mkPrepareError :: Query -> Text -> SQLError
mkPrepareError :: Query -> Text -> SQLError
mkPrepareError Query
queryText Text
msg =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage = Text
msg
        , 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
queryText
        }

mkExecuteError :: Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError :: Query -> Text -> Maybe DuckDBErrorType -> SQLError
mkExecuteError Query
queryText Text
msg Maybe DuckDBErrorType
errType =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage = Text
msg
        , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
errType
        , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Query -> Maybe Query
forall a. a -> Maybe a
Just Query
queryText
        }

throwFormatError :: Statement -> Text -> [String] -> IO a
throwFormatError :: forall a. Statement -> Text -> [String] -> IO a
throwFormatError Statement{Query
statementQuery :: Statement -> Query
statementQuery :: Query
statementQuery} Text
message [String]
params =
    FormatError -> IO a
forall e a. Exception e => e -> IO a
throwIO
        FormatError
            { formatErrorMessage :: Text
formatErrorMessage = Text
message
            , formatErrorQuery :: Query
formatErrorQuery = Query
statementQuery
            , formatErrorParams :: [String]
formatErrorParams = [String]
params
            }

throwFormatErrorBindings :: Statement -> Text -> [FieldBinding] -> IO a
throwFormatErrorBindings :: forall a. Statement -> Text -> [FieldBinding] -> IO a
throwFormatErrorBindings Statement
stmt Text
message [FieldBinding]
bindings =
    Statement -> Text -> [String] -> IO a
forall a. Statement -> Text -> [String] -> IO a
throwFormatError Statement
stmt Text
message ((FieldBinding -> String) -> [FieldBinding] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map FieldBinding -> String
renderFieldBinding [FieldBinding]
bindings)

throwFormatErrorNamed :: Statement -> Text -> [(Text, FieldBinding)] -> IO a
throwFormatErrorNamed :: forall a. Statement -> Text -> [(Text, FieldBinding)] -> IO a
throwFormatErrorNamed Statement
stmt Text
message [(Text, FieldBinding)]
bindings =
    Statement -> Text -> [String] -> IO a
forall a. Statement -> Text -> [String] -> IO a
throwFormatError Statement
stmt Text
message (((Text, FieldBinding) -> String)
-> [(Text, FieldBinding)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Text, FieldBinding) -> String
renderNamed [(Text, FieldBinding)]
bindings)
  where
    renderNamed :: (Text, FieldBinding) -> String
renderNamed (Text
name, FieldBinding
binding) =
        Text -> String
Text.unpack Text
name String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" := " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> FieldBinding -> String
renderFieldBinding FieldBinding
binding

columnIndexError :: Statement -> Int -> Maybe Int -> SQLError
columnIndexError :: Statement -> Int -> Maybe Int -> SQLError
columnIndexError Statement
stmt Int
idx Maybe Int
total =
    let base :: Text
base =
            [Text] -> Text
Text.concat
                [ String -> Text
Text.pack String
"duckdb-simple: column index "
                , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
idx)
                , String -> Text
Text.pack String
" out of bounds"
                ]
        message :: Text
message =
            case Maybe Int
total of
                Maybe Int
Nothing -> Text
base
                Just Int
count ->
                    [Text] -> Text
Text.concat
                        [ Text
base
                        , String -> Text
Text.pack String
" (column count: "
                        , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
count)
                        , String -> Text
Text.pack String
")"
                        ]
     in SQLError
            { sqlErrorMessage :: Text
sqlErrorMessage = Text
message
            , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
forall a. Maybe a
Nothing
            , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Query -> Maybe Query
forall a. a -> Maybe a
Just (Statement -> Query
statementQuery Statement
stmt)
            }

columnNameUnavailableError :: Statement -> Int -> SQLError
columnNameUnavailableError :: Statement -> Int -> SQLError
columnNameUnavailableError Statement
stmt Int
idx =
    SQLError
        { sqlErrorMessage :: Text
sqlErrorMessage =
            [Text] -> Text
Text.concat
                [ String -> Text
Text.pack String
"duckdb-simple: column name unavailable for index "
                , String -> Text
Text.pack (Int -> String
forall a. Show a => a -> String
show Int
idx)
                ]
        , sqlErrorType :: Maybe DuckDBErrorType
sqlErrorType = Maybe DuckDBErrorType
forall a. Maybe a
Nothing
        , sqlErrorQuery :: Maybe Query
sqlErrorQuery = Query -> Maybe Query
forall a. a -> Maybe a
Just (Statement -> Query
statementQuery Statement
stmt)
        }


normalizeName :: Text -> Text
normalizeName :: Text -> Text
normalizeName Text
name =
    case Text -> Maybe (Char, Text)
Text.uncons Text
name of
        Just (Char
prefix, Text
rest)
            | Char
prefix Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
':' Bool -> Bool -> Bool
|| Char
prefix Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'$' Bool -> Bool -> Bool
|| Char
prefix Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'@' -> Text
rest
        Maybe (Char, Text)
_ -> Text
name

resultRowsChanged :: Ptr DuckDBResult -> IO Int
resultRowsChanged :: Ptr DuckDBResult -> IO Int
resultRowsChanged Ptr DuckDBResult
resPtr = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr DuckDBResult -> IO DuckDBIdx
c_duckdb_rows_changed Ptr DuckDBResult
resPtr

convertRows :: (FromRow r) => Query -> [[Field]] -> IO [r]
convertRows :: forall r. FromRow r => Query -> [[Field]] -> IO [r]
convertRows = RowParser r -> Query -> [[Field]] -> IO [r]
forall r. RowParser r -> Query -> [[Field]] -> IO [r]
convertRowsWith RowParser r
forall a. FromRow a => RowParser a
fromRow

convertRowsWith :: RowParser r -> Query -> [[Field]] -> IO [r]
convertRowsWith :: forall r. RowParser r -> Query -> [[Field]] -> IO [r]
convertRowsWith RowParser r
parser Query
queryText [[Field]]
rows =
    case ([Field] -> Ok r) -> [[Field]] -> Ok [r]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (RowParser r -> [Field] -> Ok r
forall a. RowParser a -> [Field] -> Ok a
parseRow RowParser r
parser) [[Field]]
rows of
        Errors [SomeException]
err -> SQLError -> IO [r]
forall e a. Exception e => e -> IO a
throwIO (Query -> [SomeException] -> SQLError
rowErrorsToSqlError Query
queryText [SomeException]
err)
        Ok [r]
ok -> [r] -> IO [r]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [r]
ok

collectRows :: Ptr DuckDBResult -> IO [[Field]]
collectRows :: Ptr DuckDBResult -> IO [[Field]]
collectRows Ptr DuckDBResult
resPtr = do
    [StatementStreamColumn]
columns <- Ptr DuckDBResult -> IO [StatementStreamColumn]
collectResultColumns Ptr DuckDBResult
resPtr
    Int
chunkCount <- (DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr DuckDBResult -> IO DuckDBIdx
c_duckdb_result_chunk_count Ptr DuckDBResult
resPtr) :: IO Int
    if [StatementStreamColumn] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StatementStreamColumn]
columns Bool -> Bool -> Bool
|| Int
chunkCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        then [[Field]] -> IO [[Field]]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        else do
            (Int
_, [[[Field]]]
chunks) <-
                ((Int, [[[Field]]]) -> Int -> IO (Int, [[[Field]]]))
-> (Int, [[[Field]]]) -> [Int] -> IO (Int, [[[Field]]])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                    (\(Int
rowBase, [[[Field]]]
acc) Int
chunkIdx -> do
                        (Int
nextBase, [[Field]]
chunkRows) <- Ptr DuckDBResult
-> [StatementStreamColumn] -> Int -> Int -> IO (Int, [[Field]])
collectChunkRows Ptr DuckDBResult
resPtr [StatementStreamColumn]
columns Int
chunkIdx Int
rowBase
                        (Int, [[[Field]]]) -> IO (Int, [[[Field]]])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
nextBase, [[Field]]
chunkRows [[Field]] -> [[[Field]]] -> [[[Field]]]
forall a. a -> [a] -> [a]
: [[[Field]]]
acc)
                    )
                    (Int
0, [])
                    [Int
0 .. Int
chunkCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            [[Field]] -> IO [[Field]]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[[Field]]] -> [[Field]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[[Field]]] -> [[[Field]]]
forall a. [a] -> [a]
reverse [[[Field]]]
chunks))

collectResultColumns :: Ptr DuckDBResult -> IO [StatementStreamColumn]
collectResultColumns :: Ptr DuckDBResult -> IO [StatementStreamColumn]
collectResultColumns Ptr DuckDBResult
resPtr = do
    DuckDBIdx
rawCount <- Ptr DuckDBResult -> IO DuckDBIdx
c_duckdb_column_count Ptr DuckDBResult
resPtr
    let cc :: Int
cc = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
rawCount :: Int
    [Int]
-> (Int -> IO StatementStreamColumn) -> IO [StatementStreamColumn]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] \Int
columnIndex -> do
        CString
namePtr <- Ptr DuckDBResult -> DuckDBIdx -> IO CString
c_duckdb_column_name Ptr DuckDBResult
resPtr (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
columnIndex)
        Text
name <-
            if CString
namePtr 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
"column" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
columnIndex))
                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
namePtr
        DuckDBType
dtype <- Ptr DuckDBResult -> DuckDBIdx -> IO DuckDBType
c_duckdb_column_type Ptr DuckDBResult
resPtr (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
columnIndex)
        StatementStreamColumn -> IO StatementStreamColumn
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            StatementStreamColumn
                { statementStreamColumnIndex :: Int
statementStreamColumnIndex = Int
columnIndex
                , statementStreamColumnName :: Text
statementStreamColumnName = Text
name
                , statementStreamColumnType :: DuckDBType
statementStreamColumnType = DuckDBType
dtype
                }

collectChunkRows :: Ptr DuckDBResult -> [StatementStreamColumn] -> Int -> Int -> IO (Int, [[Field]])
collectChunkRows :: Ptr DuckDBResult
-> [StatementStreamColumn] -> Int -> Int -> IO (Int, [[Field]])
collectChunkRows Ptr DuckDBResult
resPtr [StatementStreamColumn]
columns Int
chunkIndex Int
rowBase = do
    DuckDBDataChunk
chunk <- Ptr DuckDBResult -> DuckDBIdx -> IO DuckDBDataChunk
c_duckdb_result_get_chunk Ptr DuckDBResult
resPtr (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
chunkIndex)
    if DuckDBDataChunk
chunk DuckDBDataChunk -> DuckDBDataChunk -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBDataChunk
forall a. Ptr a
nullPtr
        then (Int, [[Field]]) -> IO (Int, [[Field]])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
rowBase, [])
        else
            IO (Int, [[Field]]) -> IO () -> IO (Int, [[Field]])
forall a b. IO a -> IO b -> IO a
finally
                (do
                    DuckDBIdx
rawSize <- DuckDBDataChunk -> IO DuckDBIdx
c_duckdb_data_chunk_get_size DuckDBDataChunk
chunk
                    let rowCount :: Int
rowCount = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
rawSize :: Int
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
rowCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Bool -> Bool
not ([StatementStreamColumn] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StatementStreamColumn]
columns)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                        IO FieldValue -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO FieldValue -> IO ()) -> IO FieldValue -> IO ()
forall a b. (a -> b) -> a -> b
$
                            Ptr DuckDBResult -> DuckDBType -> Int -> Int -> IO FieldValue
fetchFieldValue
                                Ptr DuckDBResult
resPtr
                                (StatementStreamColumn -> DuckDBType
statementStreamColumnType ([StatementStreamColumn] -> StatementStreamColumn
forall a. HasCallStack => [a] -> a
head [StatementStreamColumn]
columns))
                                (StatementStreamColumn -> Int
statementStreamColumnIndex ([StatementStreamColumn] -> StatementStreamColumn
forall a. HasCallStack => [a] -> a
head [StatementStreamColumn]
columns))
                                Int
rowBase
                    if Int
rowCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
                        then (Int, [[Field]]) -> IO (Int, [[Field]])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
rowBase, [])
                        else do
                            [StatementStreamChunkVector]
vectors <- DuckDBDataChunk
-> [StatementStreamColumn] -> IO [StatementStreamChunkVector]
prepareChunkVectors DuckDBDataChunk
chunk [StatementStreamColumn]
columns
                            [[Field]]
rows <- (Int -> IO [Field]) -> [Int] -> IO [[Field]]
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 ([StatementStreamColumn]
-> [StatementStreamChunkVector] -> Int -> IO [Field]
buildMaterializedRow [StatementStreamColumn]
columns [StatementStreamChunkVector]
vectors) [Int
0 .. Int
rowCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                            (Int, [[Field]]) -> IO (Int, [[Field]])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
rowBase Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rowCount, [[Field]]
rows)
                )
                (DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk)

buildMaterializedRow :: [StatementStreamColumn] -> [StatementStreamChunkVector] -> Int -> IO [Field]
buildMaterializedRow :: [StatementStreamColumn]
-> [StatementStreamChunkVector] -> Int -> IO [Field]
buildMaterializedRow [StatementStreamColumn]
columns [StatementStreamChunkVector]
vectors Int
rowIdx =
    (StatementStreamColumn -> StatementStreamChunkVector -> IO Field)
-> [StatementStreamColumn]
-> [StatementStreamChunkVector]
-> IO [Field]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (Int
-> StatementStreamColumn -> StatementStreamChunkVector -> IO Field
buildMaterializedField Int
rowIdx) [StatementStreamColumn]
columns [StatementStreamChunkVector]
vectors

buildMaterializedField :: Int -> StatementStreamColumn -> StatementStreamChunkVector -> IO Field
buildMaterializedField :: Int
-> StatementStreamColumn -> StatementStreamChunkVector -> IO Field
buildMaterializedField Int
rowIdx StatementStreamColumn
column StatementStreamChunkVector{Ptr ()
statementStreamChunkVectorData :: StatementStreamChunkVector -> Ptr ()
statementStreamChunkVectorData :: Ptr ()
statementStreamChunkVectorData, Ptr DuckDBIdx
statementStreamChunkVectorValidity :: StatementStreamChunkVector -> Ptr DuckDBIdx
statementStreamChunkVectorValidity :: Ptr DuckDBIdx
statementStreamChunkVectorValidity} = do
    FieldValue
value <-
        DuckDBType -> Ptr () -> Ptr DuckDBIdx -> Int -> IO FieldValue
materializedValueFromPointers
            (StatementStreamColumn -> DuckDBType
statementStreamColumnType StatementStreamColumn
column)
            Ptr ()
statementStreamChunkVectorData
            Ptr DuckDBIdx
statementStreamChunkVectorValidity
            Int
rowIdx
    Field -> IO Field
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        Field
            { fieldName :: Text
fieldName = StatementStreamColumn -> Text
statementStreamColumnName StatementStreamColumn
column
            , fieldIndex :: Int
fieldIndex = StatementStreamColumn -> Int
statementStreamColumnIndex StatementStreamColumn
column
            , fieldValue :: FieldValue
fieldValue = FieldValue
value
            }

fetchFieldValue :: Ptr DuckDBResult -> DuckDBType -> Int -> Int -> IO FieldValue
fetchFieldValue :: Ptr DuckDBResult -> DuckDBType -> Int -> Int -> IO FieldValue
fetchFieldValue Ptr DuckDBResult
resPtr DuckDBType
dtype Int
columnIndex Int
rowIndex =
    Ptr DuckDBResult
-> Int
-> (DuckDBDataChunk -> Int -> IO FieldValue)
-> IO FieldValue
forall a.
Ptr DuckDBResult -> Int -> (DuckDBDataChunk -> Int -> IO a) -> IO a
withChunkForRow Ptr DuckDBResult
resPtr Int
rowIndex \DuckDBDataChunk
chunk Int
localRow -> do
        DuckDBVector
vector <- DuckDBDataChunk -> DuckDBIdx -> IO DuckDBVector
c_duckdb_data_chunk_get_vector DuckDBDataChunk
chunk (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
columnIndex)
        Ptr ()
dataPtr <- DuckDBVector -> IO (Ptr ())
c_duckdb_vector_get_data DuckDBVector
vector
        Ptr DuckDBIdx
validity <- DuckDBVector -> IO (Ptr DuckDBIdx)
c_duckdb_vector_get_validity DuckDBVector
vector
        DuckDBType -> Ptr () -> Ptr DuckDBIdx -> Int -> IO FieldValue
materializedValueFromPointers DuckDBType
dtype Ptr ()
dataPtr Ptr DuckDBIdx
validity Int
localRow

withChunkForRow :: Ptr DuckDBResult -> Int -> (DuckDBDataChunk -> Int -> IO a) -> IO a
withChunkForRow :: forall a.
Ptr DuckDBResult -> Int -> (DuckDBDataChunk -> Int -> IO a) -> IO a
withChunkForRow Ptr DuckDBResult
resPtr Int
targetRow DuckDBDataChunk -> Int -> IO a
action = do
    Int
chunkCount <- (DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DuckDBIdx -> Int) -> IO DuckDBIdx -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr DuckDBResult -> IO DuckDBIdx
c_duckdb_result_chunk_count Ptr DuckDBResult
resPtr) :: IO Int
    let seek :: Int -> Int -> IO a
seek Int
chunkIdx Int
remaining
            | Int
chunkIdx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
chunkCount =
                IOError -> IO a
forall e a. Exception e => e -> IO a
throwIO (String -> IOError
userError String
"duckdb-simple: row index out of bounds while materialising result")
            | Bool
otherwise = do
                DuckDBDataChunk
chunk <- Ptr DuckDBResult -> DuckDBIdx -> IO DuckDBDataChunk
c_duckdb_result_get_chunk Ptr DuckDBResult
resPtr (Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
chunkIdx)
                if DuckDBDataChunk
chunk DuckDBDataChunk -> DuckDBDataChunk -> Bool
forall a. Eq a => a -> a -> Bool
== DuckDBDataChunk
forall a. Ptr a
nullPtr
                    then Int -> Int -> IO a
seek (Int
chunkIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
remaining
                    else do
                        DuckDBIdx
rawSize <- DuckDBDataChunk -> IO DuckDBIdx
c_duckdb_data_chunk_get_size DuckDBDataChunk
chunk
                        let rowCount :: Int
rowCount = DuckDBIdx -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DuckDBIdx
rawSize :: Int
                        if Int
rowCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
                            then do
                                DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk
                                Int -> Int -> IO a
seek (Int
chunkIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
remaining
                            else if Int
remaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rowCount
                                then
                                    IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
finally
                                        (DuckDBDataChunk -> Int -> IO a
action DuckDBDataChunk
chunk Int
remaining)
                                        (DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk)
                                else do
                                    DuckDBDataChunk -> IO ()
destroyDataChunk DuckDBDataChunk
chunk
                                    Int -> Int -> IO a
seek (Int
chunkIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
remaining Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
rowCount)
    Int -> Int -> IO a
seek Int
0 Int
targetRow

materializedValueFromPointers :: DuckDBType -> Ptr () -> Ptr Word64 -> Int -> IO FieldValue
materializedValueFromPointers :: DuckDBType -> Ptr () -> Ptr DuckDBIdx -> Int -> IO FieldValue
materializedValueFromPointers DuckDBType
dtype Ptr ()
dataPtr Ptr DuckDBIdx
validity Int
rowIdx = do
    let duckIdx :: DuckDBIdx
duckIdx = Int -> DuckDBIdx
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
rowIdx :: DuckDBIdx
    Bool
valid <- Ptr DuckDBIdx -> DuckDBIdx -> IO Bool
chunkIsRowValid Ptr DuckDBIdx
validity DuckDBIdx
duckIdx
    if Bool -> Bool
not Bool
valid
        then FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FieldValue
FieldNull
        else case DuckDBType
dtype of
            DuckDBType
DuckDBTypeBoolean -> do
                Word8
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> FieldValue
FieldBool (Word8
raw Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0))
            DuckDBType
DuckDBTypeTinyInt -> do
                Int8
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int8 -> FieldValue
FieldInt8 Int8
raw)
            DuckDBType
DuckDBTypeSmallInt -> do
                Int16
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int16 -> FieldValue
FieldInt16 Int16
raw)
            DuckDBType
DuckDBTypeInteger -> do
                Int32
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int32 -> FieldValue
FieldInt32 Int32
raw)
            DuckDBType
DuckDBTypeBigInt -> do
                Int64
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int64 -> FieldValue
FieldInt64 Int64
raw)
            DuckDBType
DuckDBTypeUTinyInt -> do
                Word8
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word8 -> FieldValue
FieldWord8 Word8
raw)
            DuckDBType
DuckDBTypeUSmallInt -> do
                Word16
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word16 -> FieldValue
FieldWord16 Word16
raw)
            DuckDBType
DuckDBTypeUInteger -> do
                Word32
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word32 -> FieldValue
FieldWord32 Word32
raw)
            DuckDBType
DuckDBTypeUBigInt -> do
                DuckDBIdx
raw <- Ptr DuckDBIdx -> Int -> IO DuckDBIdx
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBIdx
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr Word64) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DuckDBIdx -> FieldValue
FieldWord64 DuckDBIdx
raw)
            DuckDBType
DuckDBTypeFloat -> do
                Float
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Float -> FieldValue
FieldFloat Float
raw)
            DuckDBType
DuckDBTypeDouble -> do
                Double
raw <- 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) Int
rowIdx
                FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> FieldValue
FieldDouble Double
raw)
            DuckDBType
DuckDBTypeVarchar -> Text -> FieldValue
FieldText (Text -> FieldValue) -> IO Text -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr () -> DuckDBIdx -> IO Text
chunkDecodeText Ptr ()
dataPtr DuckDBIdx
duckIdx
            DuckDBType
DuckDBTypeUUID -> Text -> FieldValue
FieldText (Text -> FieldValue) -> IO Text -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr () -> DuckDBIdx -> IO Text
chunkDecodeText Ptr ()
dataPtr DuckDBIdx
duckIdx
            DuckDBType
DuckDBTypeBlob -> ByteString -> FieldValue
FieldBlob (ByteString -> FieldValue) -> IO ByteString -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr () -> DuckDBIdx -> IO ByteString
chunkDecodeBlob Ptr ()
dataPtr DuckDBIdx
duckIdx
            DuckDBType
DuckDBTypeDate -> do
                Int32
raw <- 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) Int
rowIdx
                Day -> FieldValue
FieldDate (Day -> FieldValue) -> IO Day -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DuckDBDate -> IO Day
decodeDuckDBDate (Int32 -> DuckDBDate
DuckDBDate Int32
raw)
            DuckDBType
DuckDBTypeTime -> do
                DuckDBTime
raw <- Ptr DuckDBTime -> Int -> IO DuckDBTime
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBTime
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr DuckDBTime) Int
rowIdx
                TimeOfDay -> FieldValue
FieldTime (TimeOfDay -> FieldValue) -> IO TimeOfDay -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DuckDBTime -> IO TimeOfDay
decodeDuckDBTime DuckDBTime
raw
            DuckDBType
DuckDBTypeTimestamp -> do
                DuckDBTimestamp
raw <- Ptr DuckDBTimestamp -> Int -> IO DuckDBTimestamp
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBTimestamp
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr DuckDBTimestamp) Int
rowIdx
                LocalTime -> FieldValue
FieldTimestamp (LocalTime -> FieldValue) -> IO LocalTime -> IO FieldValue
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DuckDBTimestamp -> IO LocalTime
decodeDuckDBTimestamp DuckDBTimestamp
raw
            DuckDBType
DuckDBTypeDecimal -> do
                DuckDBDecimal
decimal <- Ptr DuckDBDecimal -> Int -> IO DuckDBDecimal
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff (Ptr () -> Ptr DuckDBDecimal
forall a b. Ptr a -> Ptr b
castPtr Ptr ()
dataPtr :: Ptr DuckDBDecimal) Int
rowIdx
                (Ptr DuckDBDecimal -> IO FieldValue) -> IO FieldValue
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBDecimal
ptr -> do
                    Ptr DuckDBDecimal -> DuckDBDecimal -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr DuckDBDecimal
ptr DuckDBDecimal
decimal
                    CDouble Double
value <- Ptr DuckDBDecimal -> IO CDouble
c_duckdb_decimal_to_double Ptr DuckDBDecimal
ptr
                    FieldValue -> IO FieldValue
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Double -> FieldValue
FieldDouble (Double -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Double
value))
            DuckDBType
DuckDBTypeHugeInt ->
                String -> IO FieldValue
forall a. HasCallStack => String -> a
error String
"duckdb-simple: HugeInt type not supported"
            DuckDBType
DuckDBTypeUHugeInt ->
                String -> IO FieldValue
forall a. HasCallStack => String -> a
error String
"duckdb-simple: UHugeInt type not supported"
            DuckDBType
other ->
                String -> IO FieldValue
forall a. HasCallStack => String -> a
error (String
"duckdb-simple: unsupported DuckDB type in eager result: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> DuckDBType -> String
forall a. Show a => a -> String
show DuckDBType
other)

decodeDuckDBDate :: DuckDBDate -> IO Day
decodeDuckDBDate :: DuckDBDate -> IO Day
decodeDuckDBDate DuckDBDate
raw =
    (Ptr DuckDBDateStruct -> IO Day) -> IO Day
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBDateStruct
ptr -> do
        DuckDBDate -> Ptr DuckDBDateStruct -> IO ()
c_duckdb_from_date DuckDBDate
raw Ptr DuckDBDateStruct
ptr
        DuckDBDateStruct
dateStruct <- Ptr DuckDBDateStruct -> IO DuckDBDateStruct
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBDateStruct
ptr
        Day -> IO Day
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DuckDBDateStruct -> Day
dateStructToDay DuckDBDateStruct
dateStruct)

decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay
decodeDuckDBTime :: DuckDBTime -> IO TimeOfDay
decodeDuckDBTime DuckDBTime
raw =
    (Ptr DuckDBTimeStruct -> IO TimeOfDay) -> IO TimeOfDay
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBTimeStruct
ptr -> do
        DuckDBTime -> Ptr DuckDBTimeStruct -> IO ()
c_duckdb_from_time DuckDBTime
raw Ptr DuckDBTimeStruct
ptr
        DuckDBTimeStruct
timeStruct <- Ptr DuckDBTimeStruct -> IO DuckDBTimeStruct
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBTimeStruct
ptr
        TimeOfDay -> IO TimeOfDay
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DuckDBTimeStruct -> TimeOfDay
timeStructToTimeOfDay DuckDBTimeStruct
timeStruct)

decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTime
decodeDuckDBTimestamp :: DuckDBTimestamp -> IO LocalTime
decodeDuckDBTimestamp DuckDBTimestamp
raw =
    (Ptr DuckDBTimestampStruct -> IO LocalTime) -> IO LocalTime
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca \Ptr DuckDBTimestampStruct
ptr -> do
        DuckDBTimestamp -> Ptr DuckDBTimestampStruct -> IO ()
c_duckdb_from_timestamp DuckDBTimestamp
raw Ptr DuckDBTimestampStruct
ptr
        DuckDBTimestampStruct{$sel:duckDBTimestampStructDate:DuckDBTimestampStruct :: DuckDBTimestampStruct -> DuckDBDateStruct
duckDBTimestampStructDate = DuckDBDateStruct
dateStruct, $sel:duckDBTimestampStructTime:DuckDBTimestampStruct :: DuckDBTimestampStruct -> DuckDBTimeStruct
duckDBTimestampStructTime = DuckDBTimeStruct
timeStruct} <- Ptr DuckDBTimestampStruct -> IO DuckDBTimestampStruct
forall a. Storable a => Ptr a -> IO a
peek Ptr DuckDBTimestampStruct
ptr
        LocalTime -> IO LocalTime
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            LocalTime
                { localDay :: Day
localDay = DuckDBDateStruct -> Day
dateStructToDay DuckDBDateStruct
dateStruct
                , localTimeOfDay :: TimeOfDay
localTimeOfDay = DuckDBTimeStruct -> TimeOfDay
timeStructToTimeOfDay DuckDBTimeStruct
timeStruct
                }

dateStructToDay :: DuckDBDateStruct -> Day
dateStructToDay :: DuckDBDateStruct -> Day
dateStructToDay DuckDBDateStruct{Int32
duckDBDateStructYear :: Int32
$sel:duckDBDateStructYear:DuckDBDateStruct :: DuckDBDateStruct -> Int32
duckDBDateStructYear, Int8
duckDBDateStructMonth :: Int8
$sel:duckDBDateStructMonth:DuckDBDateStruct :: DuckDBDateStruct -> Int8
duckDBDateStructMonth, Int8
duckDBDateStructDay :: Int8
$sel:duckDBDateStructDay:DuckDBDateStruct :: DuckDBDateStruct -> Int8
duckDBDateStructDay} =
    Year -> Int -> Int -> Day
fromGregorian (Int32 -> Year
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
duckDBDateStructYear) (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
duckDBDateStructMonth) (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
duckDBDateStructDay)

timeStructToTimeOfDay :: DuckDBTimeStruct -> TimeOfDay
timeStructToTimeOfDay :: DuckDBTimeStruct -> TimeOfDay
timeStructToTimeOfDay DuckDBTimeStruct{Int8
duckDBTimeStructHour :: Int8
$sel:duckDBTimeStructHour:DuckDBTimeStruct :: DuckDBTimeStruct -> Int8
duckDBTimeStructHour, Int8
duckDBTimeStructMinute :: Int8
$sel:duckDBTimeStructMinute:DuckDBTimeStruct :: DuckDBTimeStruct -> Int8
duckDBTimeStructMinute, Int8
duckDBTimeStructSecond :: Int8
$sel:duckDBTimeStructSecond:DuckDBTimeStruct :: DuckDBTimeStruct -> Int8
duckDBTimeStructSecond, Int32
duckDBTimeStructMicros :: Int32
$sel:duckDBTimeStructMicros:DuckDBTimeStruct :: DuckDBTimeStruct -> Int32
duckDBTimeStructMicros} =
    let secondsInt :: Year
secondsInt = Int8 -> Year
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
duckDBTimeStructSecond :: Integer
        micros :: Year
micros = Int32 -> Year
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
duckDBTimeStructMicros :: Integer
        fractional :: Pico
fractional = Rational -> Pico
forall a. Fractional a => Rational -> a
fromRational (Year
micros Year -> Year -> Rational
forall a. Integral a => a -> a -> Ratio a
% Year
1000000)
        totalSeconds :: Pico
totalSeconds = Year -> Pico
forall a. Num a => Year -> a
fromInteger Year
secondsInt Pico -> Pico -> Pico
forall a. Num a => a -> a -> a
+ Pico
fractional
     in Int -> Int -> Pico -> TimeOfDay
TimeOfDay
            (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
duckDBTimeStructHour)
            (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
duckDBTimeStructMinute)
            Pico
totalSeconds

peekError :: Ptr CString -> IO Text
peekError :: Ptr CString -> IO Text
peekError Ptr CString
ptr = do
    CString
errPtr <- Ptr CString -> IO CString
forall a. Storable a => Ptr a -> IO a
peek Ptr CString
ptr
    if CString
errPtr 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: failed to open database")
        else do
            String
message <- CString -> IO String
peekCString CString
errPtr
            Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Text
Text.pack String
message)

maybeFreeErr :: Ptr CString -> IO ()
maybeFreeErr :: Ptr CString -> IO ()
maybeFreeErr Ptr CString
ptr = do
    CString
errPtr <- Ptr CString -> IO CString
forall a. Storable a => Ptr a -> IO a
peek Ptr CString
ptr
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (CString
errPtr CString -> CString -> Bool
forall a. Eq a => a -> a -> Bool
/= CString
forall a. Ptr a
nullPtr) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr () -> IO ()
c_duckdb_free (CString -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr CString
errPtr)