@@ -14,11 +14,14 @@
15module Database.PostgreSQL.Simple.Internal.PQResultUtils
15module Database.PostgreSQL.Simple.Internal.PQResultUtils
16 ( finishQueryWith
16 ( finishQueryWith
17 , getRowWith
19 , getRowWith
20import Control.Exception as E
22import Control.Exception as E
21import Data.ByteString (ByteString)
23import Data.ByteString (ByteString)
24+import Data.Foldable (for_)
22import Database.PostgreSQL.Simple.FromField (ResultError(..))
25import Database.PostgreSQL.Simple.FromField (ResultError(..))
23import Database.PostgreSQL.Simple.Ok
26import Database.PostgreSQL.Simple.Ok
24import Database.PostgreSQL.Simple.Types (Query(..))
27import Database.PostgreSQL.Simple.Types (Query(..))
@@ -26,18 +29,49 @@ import Database.PostgreSQL.Simple.Internal as Base
26import Database.PostgreSQL.Simple.TypeInfo
29import Database.PostgreSQL.Simple.TypeInfo
27import qualified Database.PostgreSQL.LibPQ as PQ
30import qualified Database.PostgreSQL.LibPQ as PQ
28import qualified Data.ByteString.Char8 as B
31import qualified Data.ByteString.Char8 as B
32+import qualified Data.Vector as V
33+import qualified Data.Vector.Mutable as MV
34+import qualified Data.Vector.Unboxed as VU
35+import qualified Data.Vector.Unboxed.Mutable as MVU
29import Control.Monad.Trans.Reader
36import Control.Monad.Trans.Reader
30import Control.Monad.Trans.State.Strict
37import Control.Monad.Trans.State.Strict
32finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
39finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
33−finishQueryWith parser conn q result = do
40+finishQueryWith parser conn q result = finishQueryWith' q result $ do
41+ nrows <- PQ.ntuples result
42+ ncols <- PQ.nfields result
43+ forM' 0 (nrows-1) $ \row ->
44+ getRowWith parser row ncols conn result
46+finishQueryWithV :: RowParser r -> Connection -> Query -> PQ.Result -> IO (V.Vector r)
47+finishQueryWithV parser conn q result = finishQueryWith' q result $ do
48+ nrows <- PQ.ntuples result
49+ let PQ.Row nrows' = nrows
50+ ncols <- PQ.nfields result
51+ mv <- MV.unsafeNew (fromIntegral nrows')
52+ for_ [ 0 .. nrows-1 ] $ \row -> do
53+ let PQ.Row row' = row
54+ value <- getRowWith parser row ncols conn result
55+ MV.unsafeWrite mv (fromIntegral row') value
58+finishQueryWithVU :: VU.Unbox r => RowParser r -> Connection -> Query -> PQ.Result -> IO (VU.Vector r)
59+finishQueryWithVU parser conn q result = finishQueryWith' q result $ do
60+ nrows <- PQ.ntuples result
61+ let PQ.Row nrows' = nrows
62+ ncols <- PQ.nfields result
63+ mv <- MVU.unsafeNew (fromIntegral nrows')
64+ for_ [ 0 .. nrows-1 ] $ \row -> do
65+ let PQ.Row row' = row
66+ value <- getRowWith parser row ncols conn result
67+ MVU.unsafeWrite mv (fromIntegral row') value
70+finishQueryWith' :: Query -> PQ.Result -> IO a -> IO a
71+finishQueryWith' q result k = do
34 status <- PQ.resultStatus result
72 status <- PQ.resultStatus result
35 case status of
73 case status of
36− PQ.TuplesOk -> do
74+ PQ.TuplesOk -> k
37− nrows <- PQ.ntuples result
38− ncols <- PQ.nfields result
39− forM' 0 (nrows-1) $ \row ->
40− getRowWith parser row ncols conn result
41 PQ.EmptyQuery -> queryErr "query: Empty query"
75 PQ.EmptyQuery -> queryErr "query: Empty query"
42 PQ.CommandOk -> queryErr "query resulted in a command response"
76 PQ.CommandOk -> queryErr "query resulted in a command response"
43 PQ.CopyOut -> queryErr "query: COPY TO is not supported"
77 PQ.CopyOut -> queryErr "query: COPY TO is not supported"