@@ -69,8 +69,13 @@ module Database.PostgreSQL.Simple
6969 , query
7070 , query_
7171 -- * Queries that stream results
72+ , FetchQuantity(..)
73+ , FoldOptions
74+ , defaultFoldOptions
7275 , fold
76+ , foldWithOptions
7377 , fold_
78+ , foldWithOptions_
7479 , forEach
7580 , forEach_
7681 -- * Statements that do not return results
@@ -93,7 +98,7 @@ import Blaze.ByteString.Builder (Builder, fromByteString, toByteString)
9398import Blaze.ByteString.Builder.Char8 (fromChar)
9499import Control.Applicative ((<$>), pure)
95100import Control.Concurrent.MVar
96−import Control.Exception (Exception, bracket, onException, throw, throwIO)
101+import Control.Exception (Exception, bracket, onException, throw, throwIO, finally)
97102import Control.Monad (foldM)
98103import Control.Monad.Fix (fix)
99104import Data.ByteString (ByteString)
@@ -313,22 +318,38 @@ query_ conn q@(Query que) = do
313318-- using 'execute' instead of 'query').
314319--
315320-- * 'ResultError': result conversion failed.
316−fold :: (QueryParams q, QueryResults r) =>
317− Connection
318− -> Query -- ^ Query template.
319− -> q -- ^ Query parameters.
320− -> a -- ^ Initial state for result consumer.
321− -> (a -> r -> IO a) -- ^ Result consumer.
322− -> IO a
323−fold conn template qs z f = withConnection conn $ \c -> do
324− success <- PQ.sendQuery c =<< formatQuery conn template qs
325− if success
326− then finishFold conn c template z f
327− else do
328− msg <- maybe "fold error" id <$> PQ.errorMessage c
329− throwIO $ SqlError { sqlNativeError = -1
330− , sqlErrorMsg = msg
331− , sqlState = "" }
321+
322+fold :: ( QueryResults row, QueryParams params )
323+ => Connection
324+ -> Query
325+ -> params
326+ -> a
327+ -> (a -> row -> IO a)
328+ -> IO a
329+fold = foldWithOptions defaultFoldOptions
330+
331+data FetchQuantity
332+ = Automatic
333+ | Fixed !Int
334+
335+newtype FoldOptions
336+ = FoldOptions {
337+ fetchQuantity :: FetchQuantity
338+ }
339+
340+defaultFoldOptions = FoldOptions { fetchQuantity = Automatic }
341+
342+foldWithOptions :: ( QueryResults row, QueryParams params )
343+ => FoldOptions
344+ -> Connection
345+ -> Query
346+ -> params
347+ -> a
348+ -> (a -> row -> IO a)
349+ -> IO a
350+foldWithOptions opts conn template qs a f = do
351+ q <- formatQuery conn template qs
352+ doFold opts conn template (Query q) a f
332353
333354-- | A version of 'fold' that does not perform query substitution.
334355fold_ :: (QueryResults r) =>
@@ -337,15 +358,63 @@ fold_ :: (QueryResults r) =>
337358 -> a -- ^ Initial state for result consumer.
338359 -> (a -> r -> IO a) -- ^ Result consumer.
339360 -> IO a
340−fold_ conn q@(Query que) z f = withConnection conn $ \c -> do
341− success <- PQ.sendQuery c que
342− if success
343− then finishFold conn c q z f
344− else do
345− msg <- maybe "fold_ error" id <$> PQ.errorMessage c
346− throwIO $ SqlError { sqlNativeError = -1
347− , sqlErrorMsg = msg
348− , sqlState = "" }
361+fold_ = foldWithOptions_ defaultFoldOptions
362+
363+
364+foldWithOptions_ :: (QueryResults r) =>
365+ FoldOptions
366+ -> Connection
367+ -> Query -- ^ Query.
368+ -> a -- ^ Initial state for result consumer.
369+ -> (a -> r -> IO a) -- ^ Result consumer.
370+ -> IO a
371+foldWithOptions_ opts conn query a f = doFold opts conn query query a f
372+
373+
374+doFold :: ( QueryResults row )
375+ => FoldOptions
376+ -> Connection
377+ -> Query
378+ -> Query
379+ -> a
380+ -> (a -> row -> IO a)
381+ -> IO a
382+doFold opts conn _template q a f = do
383+ stat <- withConnection conn PQ.transactionStatus
384+ case stat of
385+ PQ.TransIdle -> withTransaction conn go
386+ PQ.TransInTrans -> go
387+ PQ.TransActive -> fail "foldWithOpts FIXME: PQ.TransActive"
388+ -- This _shouldn't_ occur in the current incarnation of
389+ -- the library, as we aren't using libpq asynchronously.
390+ -- However, it could occur in future incarnations of
391+ -- this library or if client code uses the Internal module
392+ -- to use raw libpq commands on postgresql-simple connections.
393+ PQ.TransInError -> fail "foldWithOpts FIXME: PQ.TransInError"
394+ -- This should be turned into a better error message.
395+ -- It is probably a bad idea to automatically roll
396+ -- back the transaction and start another.
397+ PQ.TransUnknown -> fail "foldWithOpts FIXME: PQ.TransUnknown"
398+ -- Not sure what this means.
399+ where
400+ go = do
401+ -- FIXME: what about name clashes with already-declared cursors?
402+ _ <- execute_ conn ("DECLARE fold NO SCROLL CURSOR FOR "
403+ `mappend` q)
404+ loop a `finally` execute_ conn "CLOSE fold"
405+
406+-- FIXME: choose the Automatic chunkSize more intelligently
407+-- One possibility is to use the type of the results, although this
408+-- still isn't a perfect solution, given that common types (e.g. text)
409+-- are of highly variable size.
410+-- A refinement of this technique is to pick this number adaptively
411+-- as results are read in from the database.
412+ chunkSize = case fetchQuantity opts of
413+ Automatic -> 256
414+ Fixed n -> n
415+ loop a = do
416+ rs <- query conn "FETCH FORWARD ? FROM fold" (Only chunkSize)
417+ if null rs then return a else foldM f a rs >>= loop
349418
350419-- | A version of 'fold' that does not transform a state value.
351420forEach :: (QueryParams q, QueryResults r) =>
@@ -404,53 +473,6 @@ finishQuery conn q result = do
404473 , sqlErrorMsg = B.concat [ "query: ", statusmsg
405474 , ": ", errormsg ]}
406475
407−
408−finishFold :: (QueryResults r) =>
409− Connection
410− -> PQ.Connection
411− -> Query -- ^ Query.
412− -> a -- ^ Initial state for result consumer.
413− -> (a -> r -> IO a) -- ^ Result consumer.
414− -> IO a
415−finishFold conn c q a_ f = loop a_
416− where
417− loop a = do
418− mres <- PQ.getResult c
419− case mres of
420− Nothing -> return a
421− Just result -> do
422− stat <- PQ.resultStatus result
423− case stat of
424− PQ.TuplesOk -> do
425− ncols <- PQ.nfields result
426− fields <- forM' 0 (ncols-1) $ \column -> do
427− type_oid <- PQ.ftype result column
428− typename <- getTypename conn type_oid
429− return Field{..}
430− nrows <- PQ.ntuples result
431− a' <- foldM (\a row -> do
432− values <- forM' 0 (ncols-1) (PQ.getvalue result row)
433− case convertResults fields values of
434− Left err -> clear c >> throwIO err
435− Right r -> f a r)
436− a [0..nrows-1]
437− loop a'
438− _ -> do
439− errormsg <- maybe "" id <$> PQ.resultErrorMessage result
440− statusmsg <- PQ.resStatus stat
441− state <- maybe "" id <$> PQ.resultErrorField result PQ.DiagSqlstate
442− clear c
443− throwIO $ SqlError {
444− sqlState = state,
445− sqlNativeError = fromEnum stat,
446− sqlErrorMsg = B.concat [ "fold: ", statusmsg
447− , ": ", errormsg ]}
448− clear c = do
449− mres <- PQ.getResult c
450− case mres of
451− Nothing -> return ()
452− Just _ -> clear c
453−
454476-- | Execute an action inside a SQL transaction.
455477--
456478-- This function initiates a transaction with a \"@begin