git.iko.soy/typeable/postgresql-simple
on master help sign in
Move cursor functionality to new Simple.Cursor module
3913d56Bardur Arantsson
postgresql-simple.cabal+1
compare with master
unifiedsplit
@@ -24,6 +24,7 @@ Library
24 Database.PostgreSQL.Simple
24 Database.PostgreSQL.Simple
25 Database.PostgreSQL.Simple.Arrays
25 Database.PostgreSQL.Simple.Arrays
26 Database.PostgreSQL.Simple.Copy
26 Database.PostgreSQL.Simple.Copy
27+ Database.PostgreSQL.Simple.Cursor
27 Database.PostgreSQL.Simple.FromField
28 Database.PostgreSQL.Simple.FromField
28 Database.PostgreSQL.Simple.FromRow
29 Database.PostgreSQL.Simple.FromRow
29 Database.PostgreSQL.Simple.LargeObjects
30 Database.PostgreSQL.Simple.LargeObjects
src/Database/PostgreSQL/Simple.hs+14−47
@@ -118,16 +118,15 @@ module Database.PostgreSQL.Simple
118 , formatQuery
118 , formatQuery
119 ) where
119 ) where
120
120
121−import Data.ByteString.Builder
121+import Data.ByteString.Builder (Builder, byteString, char8)
122− ( Builder, byteString, char8, intDec )
123import Control.Applicative ((<$>))
122import Control.Applicative ((<$>))
124import Control.Exception as E
123import Control.Exception as E
125−import Control.Monad (unless)
126import Data.ByteString (ByteString)
124import Data.ByteString (ByteString)
127import Data.Int (Int64)
125import Data.Int (Int64)
128import Data.List (intersperse)
126import Data.List (intersperse)
129import Data.Monoid (mconcat)
127import Data.Monoid (mconcat)
130−import Database.PostgreSQL.Simple.Compat ( (<>), toByteString )
128+import Database.PostgreSQL.Simple.Compat ((<>), toByteString)
129+import Database.PostgreSQL.Simple.Cursor
131import Database.PostgreSQL.Simple.FromField (ResultError(..))
130import Database.PostgreSQL.Simple.FromField (ResultError(..))
132import Database.PostgreSQL.Simple.FromRow (FromRow(..))
131import Database.PostgreSQL.Simple.FromRow (FromRow(..))
133import Database.PostgreSQL.Simple.ToField (Action(..))
132import Database.PostgreSQL.Simple.ToField (Action(..))
@@ -559,39 +558,17 @@ doFold FoldOptions{..} parser conn _template q a0 f = do
559 PQ.TransUnknown -> fail "foldWithOpts FIXME: PQ.TransUnknown"
558 PQ.TransUnknown -> fail "foldWithOpts FIXME: PQ.TransUnknown"
560 -- Not sure what this means.
559 -- Not sure what this means.
561 where
560 where
562− declare = do
561+ declare =
563− name <- newTempName conn
562+ declareCursor conn q
564− _ <- execute_ conn $ mconcat
563+ fetch cursor a =
565− [ "DECLARE ", name, " NO SCROLL CURSOR FOR ", q ]
564+ foldForwardWithParser cursor parser chunkSize f a
566− return name
565+
567− close name =
568− (execute_ conn ("CLOSE " <> name) >> return ()) `E.catch` \ex ->
569− -- Don't throw exception if CLOSE failed because the transaction is
570− -- aborted. Otherwise, it will throw away the original error.
571− unless (isFailedTransactionError ex) $ throwIO ex
572−
573− go = bracket declare close $ \(Query name) ->
566+ go = bracket declare closeCursor $ \cursor ->
574− let q = toByteString (byteString "FETCH FORWARD "
567+ let loop a = fetch cursor a >>=
575− <> intDec chunkSize
576− <> byteString " FROM "
577− <> byteString name
578− )
579− loop a = do
580− result <- exec conn q
581− status <- PQ.resultStatus result
582− case status of
568+ \r -> case r of
583− PQ.TuplesOk -> do
569+ Left a -> return a
584− nrows <- PQ.ntuples result
570+ Right a -> loop a
585− ncols <- PQ.nfields result
586− if nrows > 0
587− then do
588− let inner a row = do
589− x <- getRowWith parser row ncols conn result
590− f a x
591− foldM' inner a 0 (nrows - 1) >>= loop
592− else return a
593− _ -> throwResultError "fold" result status
594− in loop a0
571+ in loop a0
595
572
596-- FIXME: choose the Automatic chunkSize more intelligently
573-- FIXME: choose the Automatic chunkSize more intelligently
597-- One possibility is to use the type of the results, although this
574-- One possibility is to use the type of the results, although this
@@ -641,16 +618,6 @@ forEachWith_ :: RowParser r
641forEachWith_ parser conn template = foldWith_ parser conn template () . const
618forEachWith_ parser conn template = foldWith_ parser conn template () . const
642{-# INLINE forEachWith_ #-}
619{-# INLINE forEachWith_ #-}
643
620
644−foldM' :: (Ord n, Num n) => (a -> n -> IO a) -> a -> n -> n -> IO a
645−foldM' f a lo hi = loop a lo
646− where
647− loop a !n
648− | n > hi = return a
649− | otherwise = do
650− a' <- f a n
651− loop a' (n+1)
652−{-# INLINE foldM' #-}
653−
654
621
655-- $use
622-- $use
656--
623--
src/Database/PostgreSQL/Simple/Cursor.hsnew file+99
@@ -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' #-}