@@ -93,6 +93,14 @@ module Database.PostgreSQL.Simple
9393 , forEach
9494 , forEach_
9595 , returning
96+ -- ** Queries that stream results taking a parser as an argument
97+ , foldWith
98+ , foldWithOptionsAndParser
99+ , foldWith_
100+ , foldWithOptionsAndParser_
101+ , forEachWith
102+ , forEachWith_
103+ , returningWith
96104 -- * Statements that do not return results
97105 , execute
98106 , execute_
@@ -349,10 +357,13 @@ executeMany conn q qs = do
349357--
350358-- Throws 'FormatError' if the query could not be formatted correctly.
351359returning :: (ToRow q, FromRow r) => Connection -> Query -> [q] -> IO [r]
352−returning _ _ [] = return []
353−returning conn q qs = do
360+returning = returningWith fromRow
361+
362+returningWith :: (ToRow q) => RowParser r -> Connection -> Query -> [q] -> IO [r]
363+returningWith _ _ _ [] = return []
364+returningWith parser conn q qs = do
354365 result <- exec conn =<< formatMany conn q qs
355− finishQuery conn q result
366+ finishQueryWith parser conn q result
356367
357368-- | Perform a @SELECT@ or other SQL query that is expected to return
358369-- results. All results are retrieved and converted before this
@@ -427,6 +438,17 @@ fold :: ( FromRow row, ToRow params )
427438 -> IO a
428439fold = foldWithOptions defaultFoldOptions
429440
441+-- | A version of 'fold' taking a parser as an argument
442+foldWith :: ( ToRow params )
443+ => RowParser row
444+ -> Connection
445+ -> Query
446+ -> params
447+ -> a
448+ -> (a -> row -> IO a)
449+ -> IO a
450+foldWith = foldWithOptionsAndParser defaultFoldOptions
451+
430452-- | Number of rows to fetch at a time. 'Automatic' currently defaults
431453-- to 256 rows, although it might be nice to make this more intelligent
432454-- based on e.g. the average size of the rows.
@@ -461,9 +483,21 @@ foldWithOptions :: ( FromRow row, ToRow params )
461483 -> a
462484 -> (a -> row -> IO a)
463485 -> IO a
464−foldWithOptions opts conn template qs a f = do
486+foldWithOptions opts = foldWithOptionsAndParser opts fromRow
487+
488+-- | A version of 'foldWithOptions' taking a parser as an argument
489+foldWithOptionsAndParser :: (ToRow params)
490+ => FoldOptions
491+ -> RowParser row
492+ -> Connection
493+ -> Query
494+ -> params
495+ -> a
496+ -> (a -> row -> IO a)
497+ -> IO a
498+foldWithOptionsAndParser opts parser conn template qs a f = do
465499 q <- formatQuery conn template qs
466− doFold opts conn template (Query q) a f
500+ doFold opts parser conn template (Query q) a f
467501
468502-- | A version of 'fold' that does not perform query substitution.
469503fold_ :: (FromRow r) =>
@@ -474,6 +508,14 @@ fold_ :: (FromRow r) =>
474508 -> IO a
475509fold_ = foldWithOptions_ defaultFoldOptions
476510
511+-- | A version of 'foldWith' taking a parser as an argument
512+foldWith_ :: RowParser r
513+ -> Connection
514+ -> Query
515+ -> a
516+ -> (a -> r -> IO a)
517+ -> IO a
518+foldWith_ = foldWithOptionsAndParser_ defaultFoldOptions
477519
478520foldWithOptions_ :: (FromRow r) =>
479521 FoldOptions
@@ -482,18 +524,27 @@ foldWithOptions_ :: (FromRow r) =>
482524 -> a -- ^ Initial state for result consumer.
483525 -> (a -> r -> IO a) -- ^ Result consumer.
484526 -> IO a
485−foldWithOptions_ opts conn query a f = doFold opts conn query query a f
486−
487−
488−doFold :: ( FromRow row )
489− => FoldOptions
527+foldWithOptions_ opts conn query a f = doFold opts fromRow conn query query a f
528+
529+-- | A version of 'foldWithOptions_' taking a parser as an argument
530+foldWithOptionsAndParser_ :: FoldOptions
531+ -> RowParser r
532+ -> Connection
533+ -> Query -- ^ Query.
534+ -> a -- ^ Initial state for result consumer.
535+ -> (a -> r -> IO a) -- ^ Result consumer.
536+ -> IO a
537+foldWithOptionsAndParser_ opts parser conn query a f = doFold opts parser conn query query a f
538+
539+doFold :: FoldOptions
540+ -> RowParser row
490541 -> Connection
491542 -> Query
492543 -> Query
493544 -> a
494545 -> (a -> row -> IO a)
495546 -> IO a
496−doFold FoldOptions{..} conn _template q a0 f = do
547+doFold FoldOptions{..} parser conn _template q a0 f = do
497548 stat <- withConnection conn PQ.transactionStatus
498549 case stat of
499550 PQ.TransIdle -> withTransactionMode transactionMode conn go
@@ -516,7 +567,7 @@ doFold FoldOptions{..} conn _template q a0 f = do
516567 _ <- execute_ conn $ mconcat
517568 [ "DECLARE ", name, " NO SCROLL CURSOR FOR ", q ]
518569 return name
519− fetch (Query name) = query_ conn $
570+ fetch (Query name) = queryWith_ parser conn $
520571 Query (toByteString (byteString "FETCH FORWARD "
521572 <> intDec chunkSize
522573 <> byteString " FROM "
@@ -551,18 +602,37 @@ forEach :: (ToRow q, FromRow r) =>
551602 -> q -- ^ Query parameters.
552603 -> (r -> IO ()) -- ^ Result consumer.
553604 -> IO ()
554−forEach conn template qs = fold conn template qs () . const
605+forEach = forEachWith fromRow
555606{-# INLINE forEach #-}
556607
608+-- | A version of 'forEach' taking a parser as an argument
609+forEachWith :: ( ToRow q )
610+ => RowParser r
611+ -> Connection
612+ -> Query
613+ -> q
614+ -> (r -> IO ())
615+ -> IO ()
616+forEachWith parser conn template qs = foldWith parser conn template qs () . const
617+{-# INLINE forEachWith #-}
618+
557619-- | A version of 'forEach' that does not perform query substitution.
558620forEach_ :: (FromRow r) =>
559621 Connection
560622 -> Query -- ^ Query template.
561623 -> (r -> IO ()) -- ^ Result consumer.
562624 -> IO ()
563−forEach_ conn template = fold_ conn template () . const
625+forEach_ = forEachWith_ fromRow
564626{-# INLINE forEach_ #-}
565627
628+forEachWith_ :: RowParser r
629+ -> Connection
630+ -> Query
631+ -> (r -> IO ())
632+ -> IO ()
633+forEachWith_ parser conn template = foldWith_ parser conn template () . const
634+{-# INLINE forEachWith_ #-}
635+
566636forM' :: (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
567637forM' lo hi m = loop hi []
568638 where