@@ -0,0 +1,99 @@
1+{-# LANGUAGE OverloadedStrings #-}
2+
3+------------------------------------------------------------------------------
4+-- |
5+-- Module: Database.PostgreSQL.Simple.Cursor
6+-- Copyright: (c) 2011 MailRank, Inc.
7+-- (c) 2011-2012 Leon P Smith
8+-- (c) 2017 Bardur Arantsson
9+-- License: BSD3
10+-- Maintainer: Leon P Smith <leon@melding-monads.com>
11+--
12+------------------------------------------------------------------------------
13+
14+module Database.PostgreSQL.Simple.Cursor
15+ (
16+ -- * Types
17+ Cursor
18+ -- * Cursor management
19+ , declareCursor
20+ , closeCursor
21+ -- * Folding over rows from a cursor
22+ , foldForward
23+ , foldForwardWithParser
24+ ) where
25+
26+import Data.ByteString.Builder (intDec)
27+import Control.Applicative ((<$>))
28+import Control.Exception as E
29+import Control.Monad (unless, void)
30+import Data.Monoid (mconcat)
31+import Database.PostgreSQL.Simple.Compat ((<>), toByteString)
32+import Database.PostgreSQL.Simple.FromRow (FromRow(..))
33+import Database.PostgreSQL.Simple.Types (Query(..))
34+import Database.PostgreSQL.Simple.Internal as Base
35+import Database.PostgreSQL.Simple.Internal.PQResultUtils
36+import Database.PostgreSQL.Simple.Transaction
37+import qualified Database.PostgreSQL.LibPQ as PQ
38+
39+-- | Cursor within a transaction.
40+data Cursor = Cursor !Query !Connection
41+
42+-- | Declare a temporary cursor. The cursor is given a
43+-- unique name for the given connection.
44+declareCursor :: Connection -> Query -> IO Cursor
45+declareCursor conn q = do
46+ name <- newTempName conn
47+ void $ execute_ conn $ mconcat ["DECLARE ", name, " NO SCROLL CURSOR FOR ", q]
48+ return $ Cursor name conn
49+
50+-- | Close the given cursor.
51+closeCursor :: Cursor -> IO ()
52+closeCursor (Cursor name conn) =
53+ (void $ execute_ conn ("CLOSE " <> name)) `E.catch` \ex ->
54+ -- Don't throw exception if CLOSE failed because the transaction is
55+ -- aborted. Otherwise, it will throw away the original error.
56+ unless (isFailedTransactionError ex) $ throwIO ex
57+
58+-- | Fold over a chunk of rows from the given cursor, calling the
59+-- supplied fold-like function on each row as it is received. In case
60+-- the cursor is exhausted, a 'Left' value is returned, otherwise a
61+-- 'Right' value is returned.
62+foldForwardWithParser :: Cursor -> RowParser r -> Int -> (a -> r -> IO a) -> a -> IO (Either a a)
63+foldForwardWithParser (Cursor name conn) parser chunkSize f a0 = do
64+ let q = "FETCH FORWARD "
65+ <> (toByteString $ intDec chunkSize)
66+ <> " FROM "
67+ <> fromQuery name
68+ result <- exec conn q
69+ status <- PQ.resultStatus result
70+ case status of
71+ PQ.TuplesOk -> do
72+ nrows <- PQ.ntuples result
73+ ncols <- PQ.nfields result
74+ if nrows > 0
75+ then do
76+ let inner a row = do
77+ x <- getRowWith parser row ncols conn result
78+ f a x
79+ Right <$> foldM' inner a0 0 (nrows - 1)
80+ else
81+ return $ Left a0
82+ _ -> throwResultError "foldForwardWithParser" result status
83+
84+-- | Fold over a chunk of rows, calling the supplied fold-like function
85+-- on each row as it is received. In case the cursor is exhausted,
86+-- a 'Left' value is returned, otherwise a 'Right' value is returned.
87+foldForward :: FromRow r => Cursor -> Int -> (a -> r -> IO a) -> a -> IO (Either a a)
88+foldForward cursor = foldForwardWithParser cursor fromRow
89+
90+
91+foldM' :: (Ord n, Num n) => (a -> n -> IO a) -> a -> n -> n -> IO a
92+foldM' f a lo hi = loop a lo
93+ where
94+ loop a !n
95+ | n > hi = return a
96+ | otherwise = do
97+ a' <- f a n
98+ loop a' (n+1)
99+{-# INLINE foldM' #-}