Refactored Query module

src/Database/PostgreSQL/Driver.hs

@@ -20,7 +20,7 @@ module Database.PostgreSQL.Driver
     , sendBatchAndSync
     , sendBatchAndFlush
     , readNextData
-    , readReadyForQuery
+    , waitReadyForQuery
     , sendSimpleQuery
     , describeStatement
     -- * Errors

src/Database/PostgreSQL/Driver/Connection.hs

@@ -358,8 +358,8 @@ sendMessage :: RawConnection -> ClientMessage -> IO ()
 sendMessage rawConn msg = void $
     rSend rawConn . runEncode $ encodeClientMessage msg
 
-sendEncode :: RawConnection -> Encode -> IO ()
-sendEncode rawConn = void . rSend rawConn . runEncode
+sendEncode :: Connection -> Encode -> IO ()
+sendEncode conn = void . rSend (connRawConnection conn) . runEncode
 
 withConnectionMode
     :: Connection -> ConnectionMode -> (Connection -> IO a) -> IO a

src/Database/PostgreSQL/Driver/Query.hs

@@ -24,19 +24,38 @@ data Query = Query
     , qCachePolicy  :: CachePolicy
     } deriving (Show)
 
+-- | Public
+sendBatchAndFlush :: Connection -> [Query] -> IO ()
+sendBatchAndFlush = sendBatchEndBy Flush
+
 -- | Public
 sendBatchAndSync :: Connection -> [Query] -> IO ()
 sendBatchAndSync = sendBatchEndBy Sync
 
 -- | Public
-sendBatchAndFlush :: Connection -> [Query] -> IO ()
-sendBatchAndFlush = sendBatchEndBy Flush
+sendSimpleQuery :: Connection -> B.ByteString -> IO (Either Error ())
+sendSimpleQuery conn q = withConnectionMode conn SimpleQueryMode $ \c -> do
+    sendMessage (connRawConnection c) $ SimpleQuery (StatementSQL q)
+    waitReadyForQuery c
+
+-- | Public
+readNextData :: Connection -> IO (Either Error DataMessage)
+readNextData conn = readChan $ connOutDataChan conn
+
+-- | Public
+-- MUST BE called after every sended `Sync` message
+-- discards all messages preceding `ReadyForQuery`
+waitReadyForQuery :: Connection -> IO (Either Error ())
+waitReadyForQuery = fmap (>>= (liftError . findFirstError))
+                    . collectUntilReadyForQuery
+  where
+    liftError = maybe (Right ()) (Left . PostgresError)
 
 -- Helper
 sendBatchEndBy :: ClientMessage -> Connection -> [Query] -> IO ()
 sendBatchEndBy msg conn qs = do
     batch <- constructBatch conn qs
-    sendEncode (connRawConnection conn) $ batch <> encodeClientMessage msg
+    sendEncode conn $ batch <> encodeClientMessage msg
 
 constructBatch :: Connection -> [Query] -> IO Encode
 constructBatch conn = fmap fold . traverse constructSingle
@@ -65,53 +84,18 @@ constructBatch conn = fmap fold . traverse constructSingle
                 Execute pname noLimitToReceive
         pure $ parseMessage <> bindMessage <> executeMessage
 
--- | Public
-readNextData :: Connection -> IO (Either Error DataMessage)
-readNextData conn = readChan $ connOutDataChan conn
-
--- | Public
-sendSimpleQuery :: Connection -> B.ByteString -> IO (Either Error ())
-sendSimpleQuery conn q = withConnectionMode conn SimpleQueryMode $ \c -> do
-    sendMessage (connRawConnection c) $ SimpleQuery (StatementSQL q)
-    readReadyForQuery c
-
-
--- | Public
--- SHOULD BE called after every sended `Sync` message
--- skips all messages except `ReadyForQuery`
-readReadyForQuery :: Connection -> IO (Either Error ())
-readReadyForQuery = fmap (>>= (liftError . findFirstError))
-                    . collectBeforeReadyForQuery
-  where
-    liftError = maybe (Right ()) (Left . PostgresError)
-
-findFirstError :: [ServerMessage] -> Maybe ErrorDesc
-findFirstError []                       = Nothing
-findFirstError (ErrorResponse desc : _) = Just desc
-findFirstError (_ : xs)                 = findFirstError xs
-
--- Collects all messages received before ReadyForQuery
-collectBeforeReadyForQuery :: Connection -> IO (Either Error [ServerMessage])
-collectBeforeReadyForQuery conn = do
-    msg <- readChan $ connOutAllChan conn
-    case msg of
-        Left e               -> pure $ Left e
-        Right ReadForQuery{} -> pure $ Right []
-        Right m              -> fmap (m:) <$> collectBeforeReadyForQuery conn
-
 -- | Public
 describeStatement
     :: Connection
     -> B.ByteString
     -> IO (Either Error (V.Vector Oid, V.Vector FieldDescription))
 describeStatement conn stmt = do
-    sendEncode s $
+    sendEncode conn $
            encodeClientMessage (Parse sname (StatementSQL stmt) V.empty)
         <> encodeClientMessage (DescribeStatement sname)
         <> encodeClientMessage Sync
-    (parseMessages =<<) <$> collectBeforeReadyForQuery conn
+    (parseMessages =<<) <$> collectUntilReadyForQuery conn
   where
-    s = connRawConnection conn
     sname = StatementName ""
     parseMessages msgs = case msgs of
         [ParameterDescription params, NoData]
@@ -119,7 +103,21 @@ describeStatement conn stmt = do
         [ParameterDescription params, RowDescription fields]
             -> Right (params, fields)
         xs  -> Left . maybe
-             (DecodeError "Unexpected response on describe query")
+             (DecodeError "Unexpected response on a describe query")
              PostgresError
             $ findFirstError xs
 
+-- Collects all messages preceding `ReadyForQuery`
+collectUntilReadyForQuery :: Connection -> IO (Either Error [ServerMessage])
+collectUntilReadyForQuery conn = do
+    msg <- readChan $ connOutAllChan conn
+    case msg of
+        Left e               -> pure $ Left e
+        Right ReadForQuery{} -> pure $ Right []
+        Right m              -> fmap (m:) <$> collectUntilReadyForQuery conn
+
+findFirstError :: [ServerMessage] -> Maybe ErrorDesc
+findFirstError []                       = Nothing
+findFirstError (ErrorResponse desc : _) = Just desc
+findFirstError (_ : xs)                 = findFirstError xs
+

tests/Driver.hs

@@ -58,7 +58,7 @@ testBatch = withConnection $ \c -> do
     let a = "5"
         b = "3"
     sendBatchAndSync c [makeQuery1 a, makeQuery1 b]
-    readReadyForQuery c
+    waitReadyForQuery c
 
     r1 <- readNextData c
     r2 <- readNextData c
@@ -77,7 +77,7 @@ testTwoBatches = withConnection $ \c -> do
 
     sendBatchAndSync c [makeQuery2 r1 r2]
     r <- readNextData c
-    readReadyForQuery c
+    waitReadyForQuery c
 
     BS.pack (show $ a + b) @=? fromMessage r
 
@@ -93,7 +93,7 @@ testMultipleBatches = withConnection $ replicateM_ 10 . assertSingleBatch
         a @=? fromMessage r1
         r2 <- readNextData c
         b @=? fromMessage r2
-        readReadyForQuery c
+        waitReadyForQuery c
 
 -- | Query is empty string.
 testEmptyQuery :: IO ()
@@ -110,7 +110,7 @@ assertQueryNoData :: Query -> IO ()
 assertQueryNoData q = withConnection $ \c -> do
     sendBatchAndSync c [q]
     r <- fromRight <$> readNextData c
-    readReadyForQuery c
+    waitReadyForQuery c
     DataMessage [] @=? r
 
 -- | Asserts that all the received data rows are in form (Right _)
@@ -144,7 +144,7 @@ testInvalidBatch = do
   where
     assertInvalidBatch desc qs = withConnection $ \c -> do
         sendBatchAndSync c qs
-        readReadyForQuery c
+        waitReadyForQuery c
         checkInvalidResult c $ length qs
 
 -- | Describes usual statement.
@@ -189,14 +189,14 @@ testSimpleAndExtendedQuery = withConnection $ \c -> do
         b = "2"
         d = "5"
     sendBatchAndSync c [ makeQuery1 a , makeQuery1 b]
-    readReadyForQuery c
+    waitReadyForQuery c
     checkRightResult c 2
 
     rs <- sendSimpleQuery c "SELECT * FROM generate_series(1, 10)"
     assertBool "Should be Right" $ isRight rs
 
     sendBatchAndSync c [makeQuery1 d]
-    fr <- readReadyForQuery c
+    fr <- waitReadyForQuery c
     assertBool "Should be Right" $ isRight fr
     r <- fromMessage <$> readNextData c
     r @=? d
@@ -209,7 +209,7 @@ testPreparedStatementCache  = withConnection $ \c -> do
     sendBatchAndSync c [ makeQuery1 (BS.pack (show a))
                         , makeQuery1 (BS.pack (show b))
                         , makeQuery2 (BS.pack (show a)) (BS.pack (show b))]
-    readReadyForQuery c
+    waitReadyForQuery c
     r1 <- fromMessage <$> readNextData c
     r2 <- fromMessage <$> readNextData c
     r3 <- fromMessage <$> readNextData c
@@ -226,12 +226,11 @@ testPreparedStatementCache  = withConnection $ \c -> do
 testLargeQuery :: IO ()
 testLargeQuery = withConnection $ \c -> do
     sendBatchAndSync c [Query largeStmt V.empty Text Text NeverCache ]
-    readReadyForQuery c
+    waitReadyForQuery c
     r <- readNextData c
     assertBool "Should be Right" $ isRight r
   where
     largeStmt = "select typname, typnamespace, typowner, typlen, typbyval,"
                 <> "typcategory, typispreferred, typisdefined, typdelim,"
                 <> "typrelid, typelem, typarray from pg_type "
-                <> "where typtypmod = -1 and typisdefined = true"
 

tests/Protocol.hs

@@ -35,7 +35,7 @@ testSimpleQuery = withConnectionAll $ \c -> do
             <> "SELECT * FROM a;"
             <> "DROP TABLE a;"
     sendMessage rawConn $ SimpleQuery statement
-    msgs <- collectBeforeReadyForQuery c
+    msgs <- collectUntilReadyForQuery c
     assertNoErrorResponse msgs
     assertContains msgs isCommandComplete "Command complete"
   where
@@ -60,7 +60,7 @@ testExtendedQuery = withConnectionAll $ \c -> do
     sendMessage rawConn Flush
     sendMessage rawConn Sync
 
-    msgs <- collectBeforeReadyForQuery c
+    msgs <- collectUntilReadyForQuery c
     assertNoErrorResponse msgs
     assertContains msgs isBindComplete "BindComplete"
     assertContains msgs isCloseComplete "CloseComplete"
@@ -91,7 +91,7 @@ testExtendedEmptyQuery :: IO ()
 testExtendedEmptyQuery = withConnectionAll $ \c -> do
     let query = Query "" V.empty Text Text NeverCache
     sendBatchAndSync c [query]
-    msgs <- collectBeforeReadyForQuery c
+    msgs <- collectUntilReadyForQuery c
     assertNoErrorResponse msgs
     assertContains msgs isEmptyQueryResponse "EmptyQueryResponse"
   where
@@ -109,7 +109,7 @@ testExtendedQueryNoData = withConnectionAll $ \c -> do
     sendMessage rawConn $ DescribeStatement sname
     sendMessage rawConn Sync
 
-    msgs <- collectBeforeReadyForQuery c
+    msgs <- collectUntilReadyForQuery c
     assertContains msgs isNoData "NoData"
   where
     isNoData NoData = True