8000 Refactored Query module · postgres-haskell/postgres-wire@88bb3ae · GitHub
[go: up one dir, main page]
Skip to content

Commit 88bb3ae

Browse files
VyacheslavHashovVyacheslavHashov
committed
Refactored Query module
1 parent 2f877ab commit 88bb3ae

File tree

5 files changed

+55
-58
lines changed

5 files changed

+55
-58
lines changed

‎src/Database/PostgreSQL/Driver.hs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ module Database.PostgreSQL.Driver
2020
, sendBatchAndSync
2121
, sendBatchAndFlush
2222
, readNextData
23-
, readReadyForQuery
23+
, waitReadyForQuery
2424
, sendSimpleQuery
2525
, describeStatement
2626
-- * Errors

‎src/Database/PostgreSQL/Driver/Connection.hs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -358,8 +358,8 @@ sendMessage :: RawConnection -> ClientMessage -> IO ()
358358
sendMessage rawConn msg = void $
359359
rSend rawConn . runEncode $ encodeClientMessage msg
360360

361-
sendEncode :: RawConnection -> Encode -> IO ()
362-
sendEncode rawConn = void . rSend rawConn . runEncode
361+
sendEncode :: Connection -> Encode -> IO ()
362+
sendEncode conn = void . rSend (connRawConnection conn) . runEncode
363363

364364
withConnectionMode
365365
:: Connection -> ConnectionMode -> (Connection -> IO a) -> IO a

‎src/Database/PostgreSQL/Driver/Query.hs

Lines changed: 39 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -24,19 +24,38 @@ data Query = Query
2424
, qCachePolicy :: CachePolicy
2525
} deriving (Show)
2626

27+
-- | Public
28+
sendBatchAndFlush :: Connection -> [Query] -> IO ()
29+
sendBatchAndFlush = sendBatchEndBy Flush
30+
2731
-- | Public
2832
sendBatchAndSync :: Connection -> [Query] -> IO ()
2933
sendBatchAndSync = sendBatchEndBy Sync
3034

3135
-- | Public
32-
sendBatchAndFlush :: Connection -> [Query] -> IO ()
33-
sendBatchAndFlush = sendBatchEndBy Flush
36+
sendSimpleQuery :: Connection -> B.ByteString -> IO (Either Error ())
37+
sendSimpleQuery conn q = withConnectionMode conn SimpleQueryMode $ \c -> do
38+
sendMessage (connRawConnection c) $ SimpleQuery (StatementSQL q)
39+
waitReadyForQuery c
40+
41+
-- | Public
42+
readNextData :: Connection -> IO (Either Error DataMessage)
43+
readNextData conn = readChan $ connOutDataChan conn
44+
45+
-- | Public
46+
-- MUST BE called after every sended `Sync` message
47+
-- discards all messages preceding `ReadyForQuery`
48+
waitReadyForQuery :: Connection -> IO (Either Error ())
49+
waitReadyForQuery = fmap (>>= (liftError . findFirstError))
50+
. collectUntilReadyForQuery
51+
where
52+
liftError = maybe (Right ()) (Left . PostgresError)
3453

3554
-- Helper
3655
sendBatchEndBy :: ClientMessage -> Connection -> [Query] -> IO ()
3756
sendBatchEndBy msg conn qs = do
3857
batch <- constructBatch conn qs
39-
sendEncode (connRawConnection conn) $ batch <> encodeClientMessage msg
58+
sendEncode conn $ batch <> encodeClientMessage msg
4059

4160
constructBatch :: Connection -> [Query] -> IO Encode
4261
constructBatch conn = fmap fold . traverse constructSingle
@@ -65,61 +84,40 @@ constructBatch conn = fmap fold . traverse constructSingle
6584
Execute pname noLimitToReceive
6685
pure $ parseMessage <> bindMessage <> executeMessage
6786

68-
-- | Public
69-
readNextData :: Connection -> IO (Either Error DataMessage)
70-
readNextData conn = readChan $ connOutDataChan conn
71-
72-
-- | Public
73-
sendSimpleQuery :: Connection -> B.ByteString -> IO (Either Error ())
74-
sendSimpleQuery conn q = withConnectionMode conn SimpleQueryMode $ \c -> do
75-
sendMessage (connRawConnection c) $ SimpleQuery (StatementSQL q)
76-
readReadyForQuery c
77-
78-
79-
-- | Public
80-
-- SHOULD BE called after every sended `Sync` message
81-
-- skips all messages except `ReadyForQuery`
82-
readReadyForQuery :: Connection -> IO (Either Error ())
83-
readReadyForQuery = fmap (>>= (liftError . findFirstError))
84-
. collectBeforeReadyForQuery
85-
where
86-
liftError = maybe (Right ()) (Left . PostgresError)
87-
88-
findFirstError :: [ServerMessage] -> Maybe ErrorDesc
89-
findFirstError [] = Nothing
90-
findFirstError (ErrorResponse desc : _) = Just desc
91-
findFirstError (_ : xs) = findFirstError xs
92-
93-
-- Collects all messages received before ReadyForQuery
94-
collectBeforeReadyForQuery :: Connection -> IO (Either Error [ServerMessage])
95-
collectBeforeReadyForQuery conn = do
96-
msg <- readChan $ connOutAllChan conn
97-
case msg of
98-
Left e -> pure $ Left e
99-
Right ReadForQuery{} -> pure $ Right []
100-
Right m -> fmap (m:) <$> collectBeforeReadyForQuery conn
101-
10287
-- | Public
10388
describeStatement
10489
:: Connection
10590
-> B.ByteString
10691
-> IO (Either Error (V.Vector Oid, V.Vector FieldDescription))
10792
describeStatement conn stmt = do
108-
sendEncode s $
93+
sendEncode conn $
10994
encodeClientMessage (Parse sname (StatementSQL stmt) V.empty)
11095
<> encodeClientMessage (DescribeStatement sname)
11196
<> encodeClientMessage Sync
112-
(parseMessages =<<) <$> collectBeforeReadyForQuery conn
97+
(parseMessages =<<) <$> collectUntilReadyForQuery conn
11398
where
114-
s = connRawConnection conn
11599
sname = StatementName ""
116100
parseMessages msgs = case msgs of
117101
[ParameterDescription params, NoData]
118102
-> Right (params, V.empty)
119103
[ParameterDescription params, RowDescription fields]
120104
-> Right (params, fields)
121105
xs -> Left . maybe
122-
(DecodeError "Unexpected response on describe query")
106+
(DecodeError "Unexpected response on a describe query")
123107
PostgresError
124108
$ findFirstError xs
125109

110+
-- Collects all messages preceding `ReadyForQuery`
111+
collectUntilReadyForQuery :: Connection -> IO (Either Error [ServerMessage])
112+
collectUntilReadyForQuery conn = do
113+
msg <- readChan $ connOutAllChan conn
114+
case msg of
115+
Left e -> pure $ Left e
116+
Right ReadForQuery{} -> pure $ Right []
117+
Right m -> fmap (m:) <$> collectUntilReadyForQuery conn
118+
119+
findFirstError :: [ServerMessage] -> Maybe ErrorDesc
120+
findFirstError [] = Nothing
121+
findFirstError (ErrorResponse desc : _) = Just desc
122+
findFirstError (_ : xs) = findFirstError xs
123+

‎tests/Driver.hs

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ testBatch = withConnection $ \c -> do
5858
let a = "5"
5959
b = "3"
6060
sendBatchAndSync c [makeQuery1 a, makeQuery1 b]
61-
readReadyForQuery c
61+
waitReadyForQuery c
6262

6363
r1 <- readNextData c
6464
r2 <- readNextData c
@@ -77,7 +77,7 @@ testTwoBatches = withConnection $ \c -> do
7777

7878
sendBatchAndSync c [makeQuery2 r1 r2]
7979
r <- readNextData c
80-
readReadyForQuery c
80+
waitReadyForQuery c
8181

8282
BS.pack (show $ a + b) @=? fromMessage r
8383

@@ -93,7 +93,7 @@ testMultipleBatches = withConnection $ replicateM_ 10 . assertSingleBatch
9393
a @=? fromMessage r1
9494
r2 <- readNextData c
9595
b @=? fromMessage r2
96-
readReadyForQuery c
96+
waitReadyForQuery c
9797

9898
-- | Query is empty string.
9999
testEmptyQuery :: IO ()
@@ -110,7 +110,7 @@ assertQueryNoData :: Query -> IO ()
110110
assertQueryNoData q = withConnection $ \c -> do
111111
sendBatchAndSync c [q]
112112
r <- fromRight <$> readNextData c
113-
readReadyForQuery c
113+
waitReadyForQuery c
114114
DataMessage [] @=? r
115115

116116
-- | Asserts that all the received data rows are in form (Right _)
@@ -144,7 +144,7 @@ testInvalidBatch = do
144144
where
145145
assertInvalidBatch desc qs = withConnection $ \c -> do
146146
sendBatchAndSync c qs
147-
readReadyForQuery c
147+
waitReadyForQuery c
148148
checkInvalidResult c $ length qs
149149

150150
-- | Describes usual statement.
@@ -189,14 +189,14 @@ testSimpleAndExtendedQuery = withConnection $ \c -> do
189189
b = "2"
190190
d = "5"
191191
sendBatchAndSync c [ makeQuery1 a , makeQuery1 b]
192-
readReadyForQuery c
192+
waitReadyForQuery c
193193
checkRightResult c 2
194194

195195
rs <- sendSimpleQuery c "SELECT * FROM generate_series(1, 10)"
196196
assertBool "Should be Right" $ isRight rs
197197

198198
sendBatchAndSync c [makeQuery1 d]
199-
fr <- readReadyForQuery c
199+
fr <- waitReadyForQuery c
200200
assertBool "Should be Right" $ isRight fr
201201
r <- fromMessage <$> readNextData c
202202
r @=? d
@@ -209,7 +209,7 @@ testPreparedStatementCache = withConnection $ \c -> do
209209
sendBatchAndSync c [ makeQuery1 (BS.pack (show a))
210210
, makeQuery1 (BS.pack (show b))
211211
, makeQuery2 (BS.pack (show a)) (BS.pack (show b))]
212-
readReadyForQuery c
212+
waitReadyForQuery c
213213
r1 <- fromMessage <$> readNextData c
214214
r2 <- fromMessage <$> readNextData c
215215
r3 <- fromMessage <$> readNextData c
@@ -226,12 +226,11 @@ testPreparedStatementCache = withConnection $ \c -> do
226226
testLargeQuery :: IO ()
227227
testLargeQuery = withConnection $ \c -> do
228228
sendBatchAndSync c [Query largeStmt V.empty Text Text NeverCache ]
229-
readReadyForQuery c
229+
waitReadyForQuery c
230230
r <- readNextData c
231231
assertBool "Should be Right" $ isRight r
232232
where
233233
largeStmt = "select typname, typnamespace, typowner, typlen, typbyval,"
234234
<> "typcategory, typispreferred, typisdefined, typdelim,"
235235
<> "typrelid, typelem, typarray from pg_type "
236-
<> "where typtypmod = -1 and typisdefined = true"
237236

‎tests/Protocol.hs

+4Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ testSimpleQuery = withConnectionAll $ \c -> do
3535
<> "SELECT * FROM a;"
3636
<> "DROP TABLE a;"
3737
sendMessage rawConn $ SimpleQuery statement
38-
msgs <- collectBeforeReadyForQuery c
38+
msgs <- collectUntilReadyForQuery c
3939
assertNoErrorResponse msgs
4040
assertContains msgs isCommandComplete "Command complete"
4141
where
@@ -60,7 +60,7 @@ testExtendedQuery = withConnectionAll $ \c -> do
6060
sendMessage rawConn Flush
6161
sendMessage rawConn Sync
6262

63-
msgs <- collectBeforeReadyForQuery c
63+
msgs <- collectUntilReadyForQuery c
6464
assertNoErrorResponse msgs
6565
assertContains msgs isBindComplete "BindComplete"
6666
assertContains msgs isCloseComplete "CloseComplete"
@@ -91,7 +91,7 @@ testExtendedEmptyQuery :: IO ()
9191
testExtendedEmptyQuery = withConnectionAll $ \c -> do
9292
let query = Query "" V.empty Text Text NeverCache
9393
sendBatchAndSync c [query]
94-
msgs <- collectBeforeReadyForQuery c
94+
msgs <- collectUntilReadyForQuery c
9595
assertNoErrorResponse msgs
9696
assertContains msgs isEmptyQueryResponse "EmptyQueryResponse"
9797
where
@@ -109,7 +109,7 @@ testExtendedQueryNoData = withConnectionAll $ \c -> do
109109
sendMessage rawConn $ DescribeStatement sname
110110
sendMessage rawConn Sync
111111

112-
msgs <- collectBeforeReadyForQuery c
112+
msgs <- collectUntilReadyForQuery c
113113
assertContains msgs isNoData "NoData"
114114
where
115115
isNoData NoData = True

0 commit comments

Comments
 (0)
0