@@ -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 (foldM, unless, when, forM_)
126126import Data.ByteString (ByteString)
127127import Data.Int (Int64)
128128import Data.List (intersperse)
@@ -142,6 +142,8 @@ 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 qualified Control.Monad.Trans.State.Strict as SL
146+import Control.Monad.IO.Class (liftIO)
145147
146148
147149-- | Format a query string.
@@ -567,23 +569,34 @@ doFold FoldOptions{..} parser conn _template q a0 f = do
567569 _ <- execute_ conn $ mconcat
568570 [ "DECLARE ", name, " NO SCROLL CURSOR FOR ", q ]
569571 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− ))
576572 close name =
577573 (execute_ conn ("CLOSE " <> name) >> return ()) `E.catch` \ex ->
578574 -- Don't throw exception if CLOSE failed because the transaction is
579575 -- 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
586− in loop a0
576+ unless (isFailedTransactionError ex) $ throwIO ex
577+
578+ go = bracket declare close $ \(Query name) ->
579+ let q = toByteString (byteString "FETCH FORWARD "
580+ <> intDec chunkSize
581+ <> byteString " FROM "
582+ <> byteString name
583+ )
584+ loop = do
585+ result <- liftIO (exec conn q)
586+ status <- liftIO (PQ.resultStatus result)
587+ case status of
588+ PQ.TuplesOk -> do
589+ nrows <- liftIO (PQ.ntuples result)
590+ ncols <- liftIO (PQ.nfields result)
591+ when (nrows > 0) $ do
592+ forM_ [0..(nrows-1)] $ \row -> do
593+ x <- liftIO (getRowWith parser row ncols conn (Query q) result)
594+ a <- get
595+ a' <- liftIO (f a x)
596+ put a'
597+ loop
598+ _ -> liftIO (throwResultError "fold" result status)
599+ in SL.execStateT loop a0
587600
588601-- FIXME: choose the Automatic chunkSize more intelligently
589602-- One possibility is to use the type of the results, although this
@@ -651,33 +664,13 @@ finishQueryWith parser conn q result = do
651664 case status of
652665 PQ.EmptyQuery ->
653666 throwIO $ QueryError "query: Empty query" q
654− PQ.CommandOk -> do
667+ PQ.CommandOk ->
655668 throwIO $ QueryError "query resulted in a command response" q
656669 PQ.TuplesOk -> do
657− let unCol (PQ.Col x) = fromIntegral x :: Int
658670 nrows <- PQ.ntuples result
659671 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
672+ forM' 0 (nrows-1) $ \row ->
673+ getRowWith parser row ncols conn q result
681674 PQ.CopyOut ->
682675 throwIO $ QueryError "query: COPY TO is not supported" q
683676 PQ.CopyIn ->
@@ -686,6 +679,30 @@ finishQueryWith parser conn q result = do
686679 PQ.NonfatalError -> throwResultError "query" result status
687680 PQ.FatalError -> throwResultError "query" result status
688681
682+getRowWith :: RowParser r -> PQ.Row -> PQ.Column -> Connection -> Query -> PQ.Result -> IO r
683+getRowWith parser row ncols conn q result = do
684+ let rw = Row row result
685+ let unCol (PQ.Col x) = fromIntegral x :: Int
686+ okvc <- runConversion (runStateT (runReaderT (unRP parser) rw) 0) conn
687+ case okvc of
688+ Ok (val,col) | col == ncols -> return val
689+ | otherwise -> do
690+ vals <- forM' 0 (ncols-1) $ \c -> do
691+ tinfo <- getTypeInfo conn =<< PQ.ftype result c
692+ v <- PQ.getvalue result row c
693+ return ( tinfo
694+ , fmap ellipsis v )
695+ throw (ConversionFailed
696+ (show (unCol ncols) ++ " values: " ++ show vals)
697+ Nothing
698+ ""
699+ (show (unCol col) ++ " slots in target type")
700+ "mismatch between number of columns to \
701+ \convert and number in target type")
702+ Errors [] -> throwIO $ ConversionFailed "" Nothing "" "" "unknown error"
703+ Errors [x] -> throwIO x
704+ Errors xs -> throwIO $ ManyErrors xs
705+
689706ellipsis :: ByteString -> ByteString
690707ellipsis bs
691708 | B.length bs > 15 = B.take 10 bs `B.append` "[...]"