@@ -69,8 +69,6 @@ getCopyData :: Connection -> IO CopyOutResult
6969getCopyData conn = withConnection conn loop
7070 where
7171 funcName = "Database.PostgreSQL.Simple.Copy.getCopyData"
72− errCmdStatus = B.unpack funcName ++ ": failed to fetch command status"
73− errCmdStatusFmt = B.unpack funcName ++ ": failed to parse command status"
7472 loop pqconn = do
7573#if defined(mingw32_HOST_OS)
7674 row <- PQ.getCopyData pqconn False
@@ -80,12 +78,7 @@ getCopyData conn = withConnection conn loop
8078 case row of
8179 PQ.CopyOutRow rowdata -> return $! CopyOutRow rowdata
8280 PQ.CopyOutDone -> do
83− result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn
84− cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result
85− let rowCount = P.string "COPY " *> P.decimal
86− case P.parseOnly (rowCount <* P.endOfInput) cmdStat of
87− Left _ -> fail errCmdStatusFmt
88− Right n -> return $! CopyOutDone n
81+ CopyOutDone <$> getCopyCommandTag funcName pqconn
8982#if defined(mingw32_HOST_OS)
9083 PQ.CopyOutWouldBlock -> do
9184 fail (B.unpack funcName ++ ": the impossible happened")
@@ -110,26 +103,29 @@ getCopyData conn = withConnection conn loop
110103 }
111104
112105putCopyData :: Connection -> B.ByteString -> IO ()
113−putCopyData conn dat =
114− doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyData"
115− (\c -> PQ.putCopyData c dat)
116− conn
106+putCopyData conn dat = withConnection conn $ \pqconn -> do
107+ doCopyIn funcName (\c -> PQ.putCopyData c dat) pqconn
108+ where
109+ funcName = "Database.PostgreSQL.Simple.Copy.putCopyData"
110+
111+putCopyEnd :: Connection -> IO Int64
112+putCopyEnd conn = withConnection conn $ \pqconn -> do
113+ doCopyIn funcName (\c -> PQ.putCopyEnd c Nothing) pqconn
114+ getCopyCommandTag funcName pqconn
115+ where
116+ funcName = "Database.PostgreSQL.Simple.Copy.putCopyEnd"
117117
118−putCopyEnd :: Connection -> IO ()
119−putCopyEnd conn = do
120− doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyEnd"
121− (\c -> PQ.putCopyEnd c Nothing)
122− conn
118+putCopyError :: Connection -> B.ByteString -> IO Int64
119+putCopyError conn err = withConnection conn $ \pqconn -> do
120+ doCopyIn funcName (\c -> PQ.putCopyEnd c (Just err)) pqconn
121+ getCopyCommandTag funcName pqconn
122+ where
123+ funcName = "Database.PostgreSQL.Simple.Copy.putCopyError"
123124
124−putCopyError :: Connection -> B.ByteString -> IO ()
125−putCopyError conn err = do
126− doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyError"
127− (\c -> PQ.putCopyEnd c (Just err))
128− conn
129125
130126doCopyIn :: B.ByteString -> (PQ.Connection -> IO PQ.CopyInResult)
131− -> Connection -> IO ()
132−doCopyIn funcName action conn = withConnection conn loop
127+ -> PQ.Connection -> IO ()
128+doCopyIn funcName action = loop
133129 where
134130 loop pqconn = do
135131 stat <- action pqconn
@@ -152,3 +148,15 @@ doCopyIn funcName action conn = withConnection conn loop
152148 threadWaitWrite fd
153149 loop pqconn
154150{-# INLINE doCopyIn #-}
151+
152+getCopyCommandTag :: B.ByteString -> PQ.Connection -> IO Int64
153+getCopyCommandTag funcName pqconn = do
154+ result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn
155+ cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result
156+ let rowCount = P.string "COPY " *> (P.decimal <* P.endOfInput)
157+ case P.parseOnly rowCount cmdStat of
158+ Left _ -> fail errCmdStatusFmt
159+ Right n -> return $! n
160+ where
161+ errCmdStatus = B.unpack funcName ++ ": failed to fetch command status"
162+ errCmdStatusFmt = B.unpack funcName ++ ": failed to parse command status"