Merge pull request #56 from stackbuilders/all_copy_data

Add foldCopyData helper function

Author: phadej

src/Database/PostgreSQL/Simple/Copy.hs

@@ -32,6 +32,7 @@ module Database.PostgreSQL.Simple.Copy
  ( copy
  , copy_
  , CopyOutResult(..)
+ , foldCopyData
  , getCopyData
  , putCopyData
  , putCopyEnd
@@ -104,6 +105,30 @@ data CopyOutResult
  deriving (Eq, Typeable, Show)
 
 
+-- | Fold over @COPY TO STDOUT@ query passing each copied row to an accumulator
+-- and calling a post-process at the end. A connection must be in the
+-- @CopyOut@ state in order to call this function.
+--
+-- __Example__
+--
+-- > (acc, count) <- foldCopyData conn
+-- > (\acc row -> return (row:acc))
+-- > (\acc count -> return (acc, count))
+-- > []
+
+foldCopyData
+ :: Connection -- ^ Database connection
+ -> (a -> B.ByteString -> IO a) -- ^ Accumulate one row of the result
+ -> (a -> Int64 -> IO b) -- ^ Post-process accumulator with a count of rows
+ -> a -- ^ Initial accumulator
+ -> IO b -- ^ Result
+foldCopyData conn f g !acc = do
+ result <- getCopyData conn
+ case result of
+ CopyOutRow row -> f acc row >>= foldCopyData conn f g
+ CopyOutDone count -> g acc count
+
+
 -- | Retrieve some data from a @COPY TO STDOUT@ query. A connection
 -- must be in the @CopyOut@ state in order to call this function. If this
 -- returns a 'CopyOutRow', the connection remains in the @CopyOut@ state,

test/Main.hs

@@ -442,6 +442,14 @@ testCopy TestEnv{..} = do
  -- so that we can issue more queries:
  [Only (x::Int)] <- query_ conn "SELECT 2 + 2"
  x @?= 4
+ -- foldCopyData
+ copy_ conn "COPY copy_test TO STDOUT (FORMAT CSV)"
+ (acc, count) <- foldCopyData conn
+ (\acc row -> return (row:acc))
+ (\acc count -> return (acc, count))
+ []
+ sort acc @?= sort copyRows
+ count @?= 2
  where
  copyRows = ["1,foo\n"
  ,"2,bar\n"]