git.iko.soy/typeable/postgresql-simple
on master help sign in
Second attempt at implementing `fold`
6457e41Leon P Smith
src/Database/PostgreSQL/Simple.hs+95−73
compare with master
unifiedsplit
@@ -69,8 +69,13 @@ module Database.PostgreSQL.Simple
69 , query
69 , query
70 , query_
70 , query_
71 -- * Queries that stream results
71 -- * Queries that stream results
72+ , FetchQuantity(..)
73+ , FoldOptions
74+ , defaultFoldOptions
72 , fold
75 , fold
76+ , foldWithOptions
73 , fold_
77 , fold_
78+ , foldWithOptions_
74 , forEach
79 , forEach
75 , forEach_
80 , forEach_
76 -- * Statements that do not return results
81 -- * Statements that do not return results
@@ -93,7 +98,7 @@ import Blaze.ByteString.Builder (Builder, fromByteString, toByteString)
93import Blaze.ByteString.Builder.Char8 (fromChar)
98import Blaze.ByteString.Builder.Char8 (fromChar)
94import Control.Applicative ((<$>), pure)
99import Control.Applicative ((<$>), pure)
95import Control.Concurrent.MVar
100import Control.Concurrent.MVar
96−import Control.Exception (Exception, bracket, onException, throw, throwIO)
101+import Control.Exception (Exception, bracket, onException, throw, throwIO, finally)
97import Control.Monad (foldM)
102import Control.Monad (foldM)
98import Control.Monad.Fix (fix)
103import Control.Monad.Fix (fix)
99import Data.ByteString (ByteString)
104import Data.ByteString (ByteString)
@@ -313,22 +318,38 @@ query_ conn q@(Query que) = do
313-- using 'execute' instead of 'query').
318-- using 'execute' instead of 'query').
314--
319--
315-- * 'ResultError': result conversion failed.
320-- * 'ResultError': result conversion failed.
316−fold :: (QueryParams q, QueryResults r) =>
321+
322+fold :: ( QueryResults row, QueryParams params )
317− Connection
323+ => Connection
318− -> Query -- ^ Query template.
324+ -> Query
319− -> q -- ^ Query parameters.
325+ -> params
320− -> a -- ^ Initial state for result consumer.
326+ -> a
321− -> (a -> r -> IO a) -- ^ Result consumer.
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)
322− -> IO a
349+ -> IO a
323−fold conn template qs z f = withConnection conn $ \c -> do
350+foldWithOptions opts conn template qs a f = do
324− success <- PQ.sendQuery c =<< formatQuery conn template qs
351+ q <- formatQuery conn template qs
325− if success
326− then finishFold conn c template z f
352+ doFold opts conn template (Query q) a f
327− else do
328− msg <- maybe "fold error" id <$> PQ.errorMessage c
329− throwIO $ SqlError { sqlNativeError = -1
330− , sqlErrorMsg = msg
331− , sqlState = "" }
332
353
333-- | A version of 'fold' that does not perform query substitution.
354-- | A version of 'fold' that does not perform query substitution.
334fold_ :: (QueryResults r) =>
355fold_ :: (QueryResults r) =>
@@ -337,15 +358,63 @@ fold_ :: (QueryResults r) =>
337 -> a -- ^ Initial state for result consumer.
358 -> a -- ^ Initial state for result consumer.
338 -> (a -> r -> IO a) -- ^ Result consumer.
359 -> (a -> r -> IO a) -- ^ Result consumer.
339 -> IO a
360 -> IO a
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
340−fold_ conn q@(Query que) z f = withConnection conn $ \c -> do
383+ stat <- withConnection conn PQ.transactionStatus
341− success <- PQ.sendQuery c que
384+ case stat of
342− if success
385+ PQ.TransIdle -> withTransaction conn go
343− then finishFold conn c q z f
386+ PQ.TransInTrans -> go
344− else do
387+ PQ.TransActive -> fail "foldWithOpts FIXME: PQ.TransActive"
345− msg <- maybe "fold_ error" id <$> PQ.errorMessage c
388+ -- This _shouldn't_ occur in the current incarnation of
346− throwIO $ SqlError { sqlNativeError = -1
389+ -- the library, as we aren't using libpq asynchronously.
347− , sqlErrorMsg = msg
390+ -- However, it could occur in future incarnations of
348− , sqlState = "" }
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
349
418
350-- | A version of 'fold' that does not transform a state value.
419-- | A version of 'fold' that does not transform a state value.
351forEach :: (QueryParams q, QueryResults r) =>
420forEach :: (QueryParams q, QueryResults r) =>
@@ -404,53 +473,6 @@ finishQuery conn q result = do
404 , sqlErrorMsg = B.concat [ "query: ", statusmsg
473 , sqlErrorMsg = B.concat [ "query: ", statusmsg
405 , ": ", errormsg ]}
474 , ": ", errormsg ]}
406
475
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−
454-- | Execute an action inside a SQL transaction.
476-- | Execute an action inside a SQL transaction.
455--
477--
456-- This function initiates a transaction with a \"@begin
478-- This function initiates a transaction with a \"@begin