b5e7a21 · typeable/postgresql-simple

git.iko.soy/typeable/postgresql-simple

on master help sign in

Fixes and Documentation for the Copy module

b5e7a21Leon P Smith13 years ago

src/Database/PostgreSQL/Simple/Copy.hs+45−9

compare with master

unifiedsplit

@@ -1,4 +1,4 @@

1−{-# LANGUAGE CPP #-}

1+{-# LANGUAGE CPP, DeriveDataTypeable #-}

3------------------------------------------------------------------------------

4-- |

@@ -15,9 +15,9 @@

15-- To use this binding, first call 'copy' with a @COPY FROM STDIN@

16-- or @COPY TO STDOUT@ query as documented in the link above. Then

17-- call @getCopyData@ repeatedly until it returns 'CopyOutDone' in

18−-- the former case, or in the latter case call @putCopyData@ repeatedly

18+-- the former case, or in the latter, call @putCopyData@ repeatedly

19−-- and then finally call either @putCopyEnd@ to proceed or @putCopyError@

19+-- and then finish by calling either @putCopyEnd@ to proceed or

20−-- to abort.

20+-- @putCopyError@ to abort.

21--

22-- You cannot issue another query on the same connection while a copy

23-- is ongoing; this will result in an exception. It is harmless to

@@ -31,6 +31,7 @@

31module Database.PostgreSQL.Simple.Copy

32 ( copy

33 , copy_

34+ , CopyOutResult(..)

34 , getCopyData

35 , getCopyData

35 , putCopyData

36 , putCopyData

36 , putCopyEnd

37 , putCopyEnd

@@ -41,6 +42,7 @@ import Control.Applicative

41import Control.Concurrent

42import Control.Concurrent

42import Control.Exception ( throwIO )

43import Control.Exception ( throwIO )

43import qualified Data.Attoparsec.ByteString.Char8 as P

44import qualified Data.Attoparsec.ByteString.Char8 as P

45+import Data.Typeable(Typeable)

44import Data.Int(Int64)

46import Data.Int(Int64)

45import qualified Data.ByteString.Char8 as B

47import qualified Data.ByteString.Char8 as B

46import qualified Database.PostgreSQL.LibPQ as PQ

48import qualified Database.PostgreSQL.LibPQ as PQ

@@ -48,11 +50,19 @@ import Database.PostgreSQL.Simple

48import Database.PostgreSQL.Simple.Types

50import Database.PostgreSQL.Simple.Types

49import Database.PostgreSQL.Simple.Internal

51import Database.PostgreSQL.Simple.Internal

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+

51copy :: ( ToRow params ) => Connection -> Query -> params -> IO ()

57copy :: ( ToRow params ) => Connection -> Query -> params -> IO ()

52copy conn template qs = do

58copy conn template qs = do

53 q <- formatQuery conn template qs

59 q <- formatQuery conn template qs

54 doCopy "Database.PostgreSQL.Simple.Copy.copy" conn template q

60 doCopy "Database.PostgreSQL.Simple.Copy.copy" conn template q

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+

56copy_ :: Connection -> Query -> IO ()

66copy_ :: Connection -> Query -> IO ()

57copy_ conn (Query q) = do

67copy_ conn (Query q) = do

58 doCopy "Database.PostgreSQL.Simple.Copy.copy_" conn (Query q) q

68 doCopy "Database.PostgreSQL.Simple.Copy.copy_" conn (Query q) q

@@ -75,10 +85,17 @@ doCopy funcName conn template q = do

75 PQ.FatalError -> throwResultError funcName result status

85 PQ.FatalError -> throwResultError funcName result status

77data CopyOutResult

87data CopyOutResult

78− = CopyOutRow !B.ByteString -- ^ Data representing exactly one row

88+ = CopyOutRow !B.ByteString -- ^ Data representing either exactly

79− -- of the result.

89+ -- one row of the result, or header

90+ -- or footer data depending on format.

80 | CopyOutDone {-# UNPACK #-} !Int64 -- ^ No more rows, and a count of the

91 | CopyOutDone {-# UNPACK #-} !Int64 -- ^ No more rows, and a count of the

81 -- number of rows returned.

92 -- 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.

83getCopyData :: Connection -> IO CopyOutResult

100getCopyData :: Connection -> IO CopyOutResult

84getCopyData conn = withConnection conn loop

101getCopyData conn = withConnection conn loop

@@ -116,12 +133,26 @@ getCopyData conn = withConnection conn loop

116 sqlErrorHint = funcName

133 sqlErrorHint = funcName

117 }

134 }

118

135

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+

119putCopyData :: Connection -> B.ByteString -> IO ()

144putCopyData :: Connection -> B.ByteString -> IO ()

120putCopyData conn dat = withConnection conn $ \pqconn -> do

145putCopyData conn dat = withConnection conn $ \pqconn -> do

121 doCopyIn funcName (\c -> PQ.putCopyData c dat) pqconn

146 doCopyIn funcName (\c -> PQ.putCopyData c dat) pqconn

122 where

147 where

123 funcName = "Database.PostgreSQL.Simple.Copy.putCopyData"

148 funcName = "Database.PostgreSQL.Simple.Copy.putCopyData"

124

149

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+

125putCopyEnd :: Connection -> IO Int64

156putCopyEnd :: Connection -> IO Int64

126putCopyEnd conn = withConnection conn $ \pqconn -> do

157putCopyEnd conn = withConnection conn $ \pqconn -> do

127 doCopyIn funcName (\c -> PQ.putCopyEnd c Nothing) pqconn

158 doCopyIn funcName (\c -> PQ.putCopyEnd c Nothing) pqconn

@@ -129,10 +160,13 @@ putCopyEnd conn = withConnection conn $ \pqconn -> do

129 where

160 where

130 funcName = "Database.PostgreSQL.Simple.Copy.putCopyEnd"

161 funcName = "Database.PostgreSQL.Simple.Copy.putCopyEnd"

131

162

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+

132−putCopyError :: Connection -> B.ByteString -> IO Int64

167+putCopyError :: Connection -> B.ByteString -> IO ()

133putCopyError conn err = withConnection conn $ \pqconn -> do

168putCopyError conn err = withConnection conn $ \pqconn -> do

134 doCopyIn funcName (\c -> PQ.putCopyEnd c (Just err)) pqconn

169 doCopyIn funcName (\c -> PQ.putCopyEnd c (Just err)) pqconn

135− getCopyCommandTag funcName pqconn

136 where

170 where

137 funcName = "Database.PostgreSQL.Simple.Copy.putCopyError"

171 funcName = "Database.PostgreSQL.Simple.Copy.putCopyError"

138

172

@@ -167,7 +201,9 @@ getCopyCommandTag :: B.ByteString -> PQ.Connection -> IO Int64

167getCopyCommandTag funcName pqconn = do

201getCopyCommandTag funcName pqconn = do

168 result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn

202 result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn

169 cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result

203 cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result

204+ print cmdStat

170− let rowCount = P.string "COPY " *> (P.decimal <* P.endOfInput)

205+ let rowCount = P.string "COPY " *> (P.decimal <* P.endOfInput)

206+ <|> (P.endOfInput *> pure 0)

171 case P.parseOnly rowCount cmdStat of

207 case P.parseOnly rowCount cmdStat of

172 Left _ -> fail errCmdStatusFmt

208 Left _ -> fail errCmdStatusFmt

173 Right n -> return $! n

209 Right n -> return $! n