git.iko.soy/typeable/postgresql-simple
on master help sign in
Use temporary-named cursors for fold to support nesting
3a0ac85Joey Adams
src/Database/PostgreSQL/Simple.hs+19−10
compare with master
unifiedsplit
@@ -112,8 +112,7 @@ import Blaze.ByteString.Builder
112 ( Builder, fromByteString, toByteString )
112 ( Builder, fromByteString, toByteString )
113import Blaze.ByteString.Builder.Char8 (fromChar)
113import Blaze.ByteString.Builder.Char8 (fromChar)
114import Control.Applicative ((<$>), pure)
114import Control.Applicative ((<$>), pure)
115−import Control.Exception
115+import Control.Exception as E
116− ( Exception, throw, throwIO, finally )
117import Control.Monad (foldM)
116import Control.Monad (foldM)
118import Data.ByteString (ByteString)
117import Data.ByteString (ByteString)
119import Data.Int (Int64)
118import Data.Int (Int64)
@@ -463,7 +462,7 @@ doFold :: ( FromRow row )
463 -> a
462 -> a
464 -> (a -> row -> IO a)
463 -> (a -> row -> IO a)
465 -> IO a
464 -> IO a
466−doFold FoldOptions{..} conn _template q a f = do
465+doFold FoldOptions{..} conn _template q a0 f = do
467 stat <- withConnection conn PQ.transactionStatus
466 stat <- withConnection conn PQ.transactionStatus
468 case stat of
467 case stat of
469 PQ.TransIdle -> withTransactionMode transactionMode conn go
468 PQ.TransIdle -> withTransactionMode transactionMode conn go
@@ -481,10 +480,23 @@ doFold FoldOptions{..} conn _template q a f = do
481 PQ.TransUnknown -> fail "foldWithOpts FIXME: PQ.TransUnknown"
480 PQ.TransUnknown -> fail "foldWithOpts FIXME: PQ.TransUnknown"
482 -- Not sure what this means.
481 -- Not sure what this means.
483 where
482 where
484− go = do
483+ declare = do
485− -- FIXME: what about name clashes with already-declared cursors?
484+ name <- newTempName conn
486− _ <- execute_ conn ("DECLARE fold NO SCROLL CURSOR FOR " <> q)
485+ _ <- execute_ conn $ "DECLARE " <> name <> " NO SCROLL CURSOR FOR " <> q
486+ return name
487+ fetch name = query_ conn $
488+ "FETCH FORWARD " <> (Query . B.pack . show) chunkSize <> " FROM " <> name
489+ close name =
487− loop a `finally` execute_ conn "CLOSE fold"
490+ (execute_ conn ("CLOSE " <> name) >> return ()) `E.catch` \ex ->
491+ -- Don't throw exception if CLOSE failed because the transaction is
492+ -- aborted. Otherwise, it will throw away the original error.
493+ if isFailedTransactionError ex then return () else throwIO ex
494+
495+ go = bracket declare close $ \name ->
496+ let loop a = do
497+ rs <- fetch name
498+ if null rs then return a else foldM f a rs >>= loop
499+ in loop a0
488
500
489-- FIXME: choose the Automatic chunkSize more intelligently
501-- FIXME: choose the Automatic chunkSize more intelligently
490-- One possibility is to use the type of the results, although this
502-- One possibility is to use the type of the results, although this
@@ -495,9 +507,6 @@ doFold FoldOptions{..} conn _template q a f = do
495 chunkSize = case fetchQuantity of
507 chunkSize = case fetchQuantity of
496 Automatic -> 256
508 Automatic -> 256
497 Fixed n -> n
509 Fixed n -> n
498− loop a = do
499− rs <- query conn "FETCH FORWARD ? FROM fold" (Only chunkSize)
500− if null rs then return a else foldM f a rs >>= loop
501
510
502-- | A version of 'fold' that does not transform a state value.
511-- | A version of 'fold' that does not transform a state value.
503forEach :: (ToRow q, FromRow r) =>
512forEach :: (ToRow q, FromRow r) =>