@@ -122,7 +122,7 @@ import Data.ByteString.Builder
122122 ( Builder, byteString, char8, intDec )
123123import Control.Applicative ((<$>))
124124import Control.Exception as E
125−import Control.Monad (foldM)
125+import Control.Monad (unless)
126126import Data.ByteString (ByteString)
127127import Data.Int (Int64)
128128import Data.List (intersperse)
@@ -142,6 +142,7 @@ import qualified Database.PostgreSQL.LibPQ as PQ
142142import qualified Data.ByteString.Char8 as B
143143import Control.Monad.Trans.Reader
144144import Control.Monad.Trans.State.Strict
145+import Control.Monad.IO.Class (liftIO)
145146
146147
147148-- | Format a query string.
@@ -567,22 +568,37 @@ doFold FoldOptions{..} parser conn _template q a0 f = do
567568 _ <- execute_ conn $ mconcat
568569 [ "DECLARE ", name, " NO SCROLL CURSOR FOR ", q ]
569570 return name
570− fetch (Query name) = queryWith_ parser conn $
571− Query (toByteString (byteString "FETCH FORWARD "
572− <> intDec chunkSize
573− <> byteString " FROM "
574− <> byteString name
575− ))
576571 close name =
577572 (execute_ conn ("CLOSE " <> name) >> return ()) `E.catch` \ex ->
578573 -- Don't throw exception if CLOSE failed because the transaction is
579574 -- aborted. Otherwise, it will throw away the original error.
580− if isFailedTransactionError ex then return () else throwIO ex
581−
582− go = bracket declare close $ \name ->
583− let loop a = do
584− rs <- fetch name
585− if null rs then return a else foldM f a rs >>= loop
575+ unless (isFailedTransactionError ex) $ throwIO ex
576+
577+ go = bracket declare close $ \(Query name) ->
578+ let q = toByteString (byteString "FETCH FORWARD "
579+ <> intDec chunkSize
580+ <> byteString " FROM "
581+ <> byteString name
582+ )
583+ loop a = do
584+ result <- liftIO (exec conn q)
585+ status <- liftIO (PQ.resultStatus result)
586+ case status of
587+ PQ.TuplesOk -> do
588+ nrows <- liftIO (PQ.ntuples result)
589+ ncols <- liftIO (PQ.nfields result)
590+ if nrows > 0
591+ then do
592+ a' <- innerLoop 0 nrows ncols result a
593+ loop a'
594+ else return a
595+ _ -> liftIO (throwResultError "fold" result status)
596+ innerLoop row nrows ncols result a
597+ | row < nrows = do
598+ x <- getRowWith parser row ncols conn result
599+ a' <- f a x
600+ innerLoop (row+1) nrows ncols result a'
601+ | otherwise = return a
586602 in loop a0
587603
588604-- FIXME: choose the Automatic chunkSize more intelligently
@@ -651,33 +667,13 @@ finishQueryWith parser conn q result = do
651667 case status of
652668 PQ.EmptyQuery ->
653669 throwIO $ QueryError "query: Empty query" q
654− PQ.CommandOk -> do
670+ PQ.CommandOk ->
655671 throwIO $ QueryError "query resulted in a command response" q
656672 PQ.TuplesOk -> do
657− let unCol (PQ.Col x) = fromIntegral x :: Int
658673 nrows <- PQ.ntuples result
659674 ncols <- PQ.nfields result
660− forM' 0 (nrows-1) $ \row -> do
661− let rw = Row row result
662− okvc <- runConversion (runStateT (runReaderT (unRP parser) rw) 0) conn
663− case okvc of
664− Ok (val,col) | col == ncols -> return val
665− | otherwise -> do
666− vals <- forM' 0 (ncols-1) $ \c -> do
667− tinfo <- getTypeInfo conn =<< PQ.ftype result c
668− v <- PQ.getvalue result row c
669− return ( tinfo
670− , fmap ellipsis v )
671− throw (ConversionFailed
672− (show (unCol ncols) ++ " values: " ++ show vals)
673− Nothing
674− ""
675− (show (unCol col) ++ " slots in target type")
676− "mismatch between number of columns to \
677− \convert and number in target type")
678− Errors [] -> throwIO $ ConversionFailed "" Nothing "" "" "unknown error"
679− Errors [x] -> throwIO x
680− Errors xs -> throwIO $ ManyErrors xs
675+ forM' 0 (nrows-1) $ \row ->
676+ getRowWith parser row ncols conn result
681677 PQ.CopyOut ->
682678 throwIO $ QueryError "query: COPY TO is not supported" q
683679 PQ.CopyIn ->
@@ -686,6 +682,30 @@ finishQueryWith parser conn q result = do
686682 PQ.NonfatalError -> throwResultError "query" result status
687683 PQ.FatalError -> throwResultError "query" result status
688684
685+getRowWith :: RowParser r -> PQ.Row -> PQ.Column -> Connection -> PQ.Result -> IO r
686+getRowWith parser row ncols conn result = do
687+ let rw = Row row result
688+ let unCol (PQ.Col x) = fromIntegral x :: Int
689+ okvc <- runConversion (runStateT (runReaderT (unRP parser) rw) 0) conn
690+ case okvc of
691+ Ok (val,col) | col == ncols -> return val
692+ | otherwise -> do
693+ vals <- forM' 0 (ncols-1) $ \c -> do
694+ tinfo <- getTypeInfo conn =<< PQ.ftype result c
695+ v <- PQ.getvalue result row c
696+ return ( tinfo
697+ , fmap ellipsis v )
698+ throw (ConversionFailed
699+ (show (unCol ncols) ++ " values: " ++ show vals)
700+ Nothing
701+ ""
702+ (show (unCol col) ++ " slots in target type")
703+ "mismatch between number of columns to \
704+ \convert and number in target type")
705+ Errors [] -> throwIO $ ConversionFailed "" Nothing "" "" "unknown error"
706+ Errors [x] -> throwIO x
707+ Errors xs -> throwIO $ ManyErrors xs
708+
689709ellipsis :: ByteString -> ByteString
690710ellipsis bs
691711 | B.length bs > 15 = B.take 10 bs `B.append` "[...]"