@@ -0,0 +1,155 @@
1+{-# LANGUAGE CPP #-}
2+
3+------------------------------------------------------------------------------
4+-- |
5+-- Module: Database.PostgreSQL.Simple.Copy
6+-- Copyright: (c) 2013 Leon P Smith
7+-- License: BSD3
8+-- Maintainer: Leon P Smith <leon@melding-monads.com>
9+-- Stability: experimental
10+--
11+-- mid-level support for COPY IN and COPY OUT.
12+--
13+------------------------------------------------------------------------------
14+
15+module Database.PostgreSQL.Simple.Copy
16+ ( copy
17+ , copy_
18+ , getCopyData
19+ , putCopyData
20+ , putCopyEnd
21+ , putCopyError
22+ ) where
23+
24+import Control.Applicative
25+import Control.Concurrent ( threadWaitRead, threadWaitWrite )
26+import Control.Exception ( throwIO )
27+import qualified Data.Attoparsec.ByteString.Char8 as P
28+import Data.Int(Int64)
29+import qualified Data.ByteString.Char8 as B
30+import qualified Database.PostgreSQL.LibPQ as PQ
31+import Database.PostgreSQL.Simple hiding
32+ ( fold, fold_, forEach, forEach_ )
33+import Database.PostgreSQL.Simple.Types
34+import Database.PostgreSQL.Simple.Internal
35+
36+copy :: ( ToRow params ) => Connection -> Query -> params -> IO ()
37+copy conn template qs = do
38+ q <- formatQuery conn template qs
39+ doCopy "Database.PostgreSQL.Simple.Copy.copy" conn template q
40+
41+copy_ :: Connection -> Query -> IO ()
42+copy_ conn (Query q) = do
43+ doCopy "Database.PostgreSQL.Simple.Copy.copy_" conn (Query q) q
44+
45+doCopy :: B.ByteString -> Connection -> Query -> B.ByteString -> IO ()
46+doCopy funcName conn template q = do
47+ result <- exec conn q
48+ status <- PQ.resultStatus result
49+ let err = throwIO $ QueryError
50+ (B.unpack funcName ++ " " ++ show status)
51+ template
52+ case status of
53+ PQ.EmptyQuery -> err
54+ PQ.CommandOk -> err
55+ PQ.TuplesOk -> err
56+ PQ.CopyOut -> return ()
57+ PQ.CopyIn -> return ()
58+ PQ.BadResponse -> throwResultError funcName result status
59+ PQ.NonfatalError -> throwResultError funcName result status
60+ PQ.FatalError -> throwResultError funcName result status
61+{-# INLINE doCopy #-}
62+
63+data CopyOutResult
64+ = CopyOutRow !B.ByteString -- ^ Data representing exactly one row
65+ -- of the result.
66+ | CopyOutDone {-# UNPACK #-} !Int64 -- ^ No more rows, and a count of the
67+ -- number of rows returned.
68+
69+getCopyData :: Connection -> IO CopyOutResult
70+getCopyData conn = withConnection conn loop
71+ where
72+ funcName = "Database.PostgreSQL.Simple.Copy.getCopyData"
73+ errCmdStatus = B.unpack funcName ++ ": failed to fetch command status"
74+ errCmdStatusFmt = B.unpack funcName ++ ": failed to parse command status"
75+ loop pqconn = do
76+#if defined(mingw32_HOST_OS)
77+ row <- PQ.getCopyData pqconn False
78+#else
79+ row <- PQ.getCopyData pqconn True
80+#endif
81+ case row of
82+ PQ.CopyOutRow rowdata -> return (CopyOutRow rowdata)
83+ PQ.CopyOutDone -> do
84+ result <- maybe (fail errCmdStatus) return =<< PQ.getResult pqconn
85+ cmdStat <- maybe (fail errCmdStatus) return =<< PQ.cmdStatus result
86+ let rowCount = P.string "COPY " *> P.decimal
87+ case P.parseOnly rowCount cmdStat of
88+ Left _ -> fail errCmdStatusFmt
89+ Right n -> return (CopyOutDone n)
90+#if defined(mingw32_HOST_OS)
91+ PQ.CopyOutWouldBlock -> do
92+ fail (B.unpack funcName ++ ": the impossible happened")
93+#else
94+ PQ.CopyOutWouldBlock -> do
95+ mfd <- PQ.socket pqconn
96+ case mfd of
97+ Nothing -> throwIO (fdError funcName)
98+ Just fd -> do
99+ threadWaitRead fd
100+ _ <- PQ.consumeInput pqconn
101+ loop pqconn
102+#endif
103+ PQ.CopyOutError -> do
104+ mmsg <- PQ.errorMessage pqconn
105+ throwIO SqlError {
106+ sqlState = "",
107+ sqlExecStatus = FatalError,
108+ sqlErrorMsg = maybe "" id mmsg,
109+ sqlErrorDetail = "",
110+ sqlErrorHint = funcName
111+ }
112+
113+putCopyData :: Connection -> B.ByteString -> IO ()
114+putCopyData conn dat =
115+ doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyData"
116+ (\c -> PQ.putCopyData c dat)
117+ conn
118+
119+putCopyEnd :: Connection -> IO ()
120+putCopyEnd conn = do
121+ doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyEnd"
122+ (\c -> PQ.putCopyEnd c Nothing)
123+ conn
124+
125+putCopyError :: Connection -> B.ByteString -> IO ()
126+putCopyError conn err = do
127+ doCopyIn "Database.PostgreSQL.Simple.Copy.putCopyError"
128+ (\c -> PQ.putCopyEnd c (Just err))
129+ conn
130+
131+doCopyIn :: B.ByteString -> (PQ.Connection -> IO PQ.CopyInResult)
132+ -> Connection -> IO ()
133+doCopyIn funcName action conn = withConnection conn loop
134+ where
135+ loop pqconn = do
136+ stat <- action pqconn
137+ case stat of
138+ PQ.CopyInOk -> return ()
139+ PQ.CopyInError -> do
140+ mmsg <- PQ.errorMessage pqconn
141+ throwIO SqlError {
142+ sqlState = "",
143+ sqlExecStatus = FatalError,
144+ sqlErrorMsg = maybe "" id mmsg,
145+ sqlErrorDetail = "",
146+ sqlErrorHint = funcName
147+ }
148+ PQ.CopyInWouldBlock -> do
149+ mfd <- PQ.socket pqconn
150+ case mfd of
151+ Nothing -> throwIO (fdError funcName)
152+ Just fd -> do
153+ threadWaitWrite fd
154+ loop pqconn
155+{-# INLINE doCopyIn #-}