expose PQsetChunkedRowsMode

src/Database/PostgreSQL/LibPQ.hs

@@ -168,6 +168,7 @@ module Database.PostgreSQL.LibPQ
     , setnonblocking
     , isnonblocking
     , setSingleRowMode
+    , setChunkedRowsMode
     , FlushStatus(..)
     , flush
 
@@ -1619,17 +1620,35 @@ isnonblocking connection = enumFromConn connection c_PQisnonblocking
 
 -- | Select single-row mode for the currently-executing query.
 --
--- This function can only be called immediately after PQsendQuery or one of its
+-- This function can only be called immediately after 'sendQuery' or one of its
 -- sibling functions, before any other operation on the connection such as
--- PQconsumeInput or PQgetResult. If called at the correct time, the function
--- activates single-row mode for the current query and returns 1. Otherwise the
--- mode stays unchanged and the function returns 0. In any case, the mode
--- reverts to normal after completion of the current query.
+-- 'consumeInput' or 'getResult'. If called at the correct time, the function
+-- activates single-row mode for the current query and returns 'True'.
+-- Otherwise the mode stays unchanged and the function returns 'False'. In any
+-- case, the mode reverts to normal after completion of the current query.
 setSingleRowMode :: Connection
                  -> IO Bool
 setSingleRowMode connection = enumFromConn connection c_PQsetSingleRowMode
 
 
+-- | Select chunked mode for the currently-executing query.
+--
+-- This function is similar to 'setSingleRowMode', except that it specifies
+-- retrieval of up to @chunkSize@ rows per 'Result', not necessarily just one
+-- row. This function can only be called immediately after 'sendQuery' or one
+-- of its sibling functions, before any other operation on the connection such
+-- as 'consumeInput' or 'getResult'. If called at the correct time, the
+-- function activates chunked mode for the current query and returns 'True'.
+-- Otherwise the mode stays unchanged and the function returns 'False'. In any
+-- case, the mode reverts to normal after completion of the current query.
+setChunkedRowsMode :: Connection
+                   -> Int
+                   -> IO Bool
+setChunkedRowsMode connection chunkSize =
+    enumFromConn connection $ \p ->
+        c_PQsetChunkedRowsMode p (fromIntegral chunkSize)
+
+
 data FlushStatus = FlushOk
                  | FlushFailed
                  | FlushWriting

src/Database/PostgreSQL/LibPQ/FFI.hs

@@ -118,7 +118,7 @@ foreign import capi        "hs-libpq.h PQputCopyData"
 
 foreign import capi        "hs-libpq.h PQputCopyEnd"
     c_PQputCopyEnd :: Ptr PGconn -> CString -> IO CInt
- 
+
 -- TODO: GHC #22043
 foreign import ccall       "hs-libpq.h PQgetCopyData"
     c_PQgetCopyData :: Ptr PGconn -> Ptr (Ptr Word8) -> CInt -> IO CInt
@@ -177,6 +177,9 @@ foreign import capi        "hs-libpq.h PQisnonblocking"
 foreign import capi        "hs-libpq.h PQsetSingleRowMode"
     c_PQsetSingleRowMode :: Ptr PGconn -> IO CInt
 
+foreign import capi        "hs-libpq.h PQsetChunkedRowsMode"
+    c_PQsetChunkedRowsMode :: Ptr PGconn -> CInt -> IO CInt
+
 foreign import capi        "hs-libpq.h PQgetResult"
     c_PQgetResult :: Ptr PGconn -> IO (Ptr PGresult)
 

test/Smoke.hs

@@ -18,6 +18,7 @@ main = do
         [ testCaseSteps "smoke" $ smoke connString
         , testCaseSteps "issue54" $ issue54 connString
         , testCaseSteps "pipeline" $ testPipeline connString
+        , testCaseSteps "rowmodes" $ testRowModes connString
         ]
 
 withConnstring :: (BS8.ByteString -> IO ()) -> IO ()
@@ -114,8 +115,43 @@ testPipeline connstring info = do
 
     finish conn
   where
-    shouldBe r value = assertEqual "shouldBe" r value
+    shouldBe r value = assertEqual "shouldBe" value r
 
     shouldReturn action value = do
         r <- action
         r `shouldBe` value
+
+testRowModes :: BS8.ByteString -> (String -> IO ()) -> IO ()
+testRowModes connstring info = do
+    conn <- connectdb connstring
+
+    let q = sendQuery conn (BS8.pack "select * from (values (1), (2), (3))") >>= assertEqual "sendQuery" True
+
+    do q
+       Just r1 <- getResult conn
+       ntuples r1 >>= assertEqual "no row mode" 3
+       getResult conn >>= assertEqual "no row mode/end" Nothing
+
+    do q
+       setSingleRowMode conn >>= assertEqual "setSingleRowMode" True
+       Just r1 <- getResult conn
+       ntuples r1 >>= assertEqual "singlerow 1" 1
+       Just r2 <- getResult conn
+       ntuples r2 >>= assertEqual "singlerow 2" 1
+       Just r3 <- getResult conn
+       ntuples r3 >>= assertEqual "singlerow 3" 1
+       Just r4 <- getResult conn
+       ntuples r4 >>= assertEqual "singlerow eof" 0
+       getResult conn >>= assertEqual "singlerow end" Nothing
+
+    do q
+       setChunkedRowsMode conn 2 >>= assertEqual "setChunkedRowsMode" True
+       Just r1 <- getResult conn
+       ntuples r1 >>= assertEqual "chunkedrow 1" 2
+       Just r2 <- getResult conn
+       ntuples r2 >>= assertEqual "chunkedrow 2" 1
+       Just r3 <- getResult conn
+       ntuples r3 >>= assertEqual "chunkedrow eof" 0
+       getResult conn >>= assertEqual "chunkedrow end" Nothing
+
+    finish conn