@@ -1,4 +1,4 @@
1−{-# LANGUAGE CPP #-}
1+{-# LANGUAGE CPP, DeriveDataTypeable #-}
22
33------------------------------------------------------------------------------
44-- |
@@ -15,9 +15,9 @@
1515-- To use this binding, first call 'copy' with a @COPY FROM STDIN@
1616-- or @COPY TO STDOUT@ query as documented in the link above. Then
1717-- call @getCopyData@ repeatedly until it returns 'CopyOutDone' in
18−-- the former case, or in the latter case call @putCopyData@ repeatedly
19−-- and then finally call either @putCopyEnd@ to proceed or @putCopyError@
20−-- to abort.
18+-- the former case, or in the latter, call @putCopyData@ repeatedly
19+-- and then finish by calling either @putCopyEnd@ to proceed or
20+-- @putCopyError@ to abort.
2121--
2222-- You cannot issue another query on the same connection while a copy
2323-- is ongoing; this will result in an exception. It is harmless to
@@ -31,6 +31,7 @@
3131module Database.PostgreSQL.Simple.Copy
3232 ( copy
3333 , copy_
34+ , CopyOutResult(..)
3435 , getCopyData
3536 , putCopyData
3637 , putCopyEnd
@@ -41,6 +42,7 @@ import Control.Applicative
4142import Control.Concurrent
4243import Control.Exception ( throwIO )
4344import qualified Data.Attoparsec.ByteString.Char8 as P
45+import Data.Typeable(Typeable)
4446import Data.Int(Int64)
4547import qualified Data.ByteString.Char8 as B
4648import qualified Database.PostgreSQL.LibPQ as PQ
@@ -48,11 +50,19 @@ import Database.PostgreSQL.Simple
4850import Database.PostgreSQL.Simple.Types
4951import Database.PostgreSQL.Simple.Internal
5052
53+-- | Issue a query that changes a connection's state to @CopyIn@
54+-- (via a @COPY FROM STDIN@ query) or @CopyOut@ (via @COPY TO STDOUT@)
55+-- query. Performs parameter subsitution.
56+
5157copy :: ( ToRow params ) => Connection -> Query -> params -> IO ()
5258copy conn template qs = do
5359 q <- formatQuery conn template qs
5460 doCopy "Database.PostgreSQL.Simple.Copy.copy" conn template q
5561
62+-- | Issue a query that changes a connection's state to @CopyIn@
63+-- (via a @COPY FROM STDIN@ query) or @CopyOut@ (via @COPY TO STDOUT@)
64+-- query. Does not perform parameter subsitution.
65+
5666copy_ :: Connection -> Query -> IO ()
5767copy_ conn (Query q) = do
5868 doCopy "Database.PostgreSQL.Simple.Copy.copy_" conn (Query q) q
@@ -75,10 +85,17 @@ doCopy funcName conn template q = do
7585 PQ.FatalError -> throwResultError funcName result status
7686
7787data CopyOutResult
78− = CopyOutRow !B.ByteString -- ^ Data representing exactly one row
79− -- of the result.
88+ = CopyOutRow !B.ByteString -- ^ Data representing either exactly
89+ -- one row of the result, or header
90+ -- or footer data depending on format.
8091 | CopyOutDone {-# UNPACK #-} !Int64 -- ^ No more rows, and a count of the
8192 -- number of rows returned.
93+ deriving (Eq, Typeable, Show)
94+
95+-- | A connection must be in the @CopyOut@ state in order to call this
96+-- function, via a @COPY TO STDOUT@ query. If this returns a 'CopyOutRow',
97+-- the connection remains in the @CopyOut@ state, if it returns 'CopyOutDone',
98+-- then the connection has reverted to the ready state.
8299
83100getCopyData :: Connection -> IO CopyOutResult
84101getCopyData conn = withConnection conn loop
@@ -116,12 +133,26 @@ getCopyData conn = withConnection conn loop
116133 sqlErrorHint = funcName
117134 }
118135
136+-- | A connection must be in the @CopyIn@ state in order to call this
137+-- function, via a @COPY FROM STDIN@ query. The connection remains
138+-- in a @CopyIn@ state after this function is called. Note that
139+-- the data does not need to represent a single row, or even an
140+-- integral number of rows. The net result of
141+-- @putCopyData conn a >> putCopyData conn b@
142+-- is the same as @putCopyData conn c@ whenever @c == BS.append a b@.
143+
119144putCopyData :: Connection -> B.ByteString -> IO ()
120145putCopyData conn dat = withConnection conn $ \pqconn -> do
121146 doCopyIn funcName (\c -> PQ.putCopyData c dat) pqconn
122147 where
123148 funcName = "Database.PostgreSQL.Simple.Copy.putCopyData"
124149
150+
151+-- | A connection must be in the @CopyIn@ state in order to call this
152+-- function, via a @COPY FROM STDIN@ query. Completes the COPY IN
153+-- operation, changing the connection's state back to normal.
154+-- Returns the number of rows processed.
155+
125156putCopyEnd :: Connection -> IO Int64
126157putCopyEnd conn = withConnection conn $ \pqconn -> do
127158 doCopyIn funcName (\c -> PQ.putCopyEnd c Nothing) pqconn
@@ -129,10 +160,13 @@ putCopyEnd conn = withConnection conn $ \pqconn -> do
129160 where
130161 funcName = "Database.PostgreSQL.Simple.Copy.putCopyEnd"
131162
132−putCopyError :: Connection -> B.ByteString -> IO Int64
163+-- | A connection must be in the @CopyIn@ state in order to call this
164+-- function, via a @COPY FROM STDIN@ query. Aborts the COPY IN
165+-- operation, changing the connection's state back to normal.
166+
167+putCopyError :: Connection -> B.ByteString -> IO ()
133168putCopyError conn err = withConnection conn $ \pqconn -> do
134169 doCopyIn funcName (\c -> PQ.putCopyEnd c (Just err)) pqconn
135− getCopyCommandTag funcName pqconn
136170 where
137171 funcName = "Database.PostgreSQL.Simple.Copy.putCopyError"
138172
@@ -167,7 +201,9 @@ getCopyCommandTag :: B.ByteString -> PQ.Connection -> IO Int64
167201getCopyCommandTag funcName pqconn = do
168202 result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn
169203 cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result
170− let rowCount = P.string "COPY " *> (P.decimal <* P.endOfInput)
204+ print cmdStat
205+ let rowCount = P.string "COPY " *> (P.decimal <* P.endOfInput)
206+ <|> (P.endOfInput *> pure 0)
171207 case P.parseOnly rowCount cmdStat of
172208 Left _ -> fail errCmdStatusFmt
173209 Right n -> return $! n