diff --git a/src/Database/PostgreSQL/LibPQ.hs b/src/Database/PostgreSQL/LibPQ.hs index 5b0c625..36e3222 100644 --- a/src/Database/PostgreSQL/LibPQ.hs +++ b/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 diff --git a/src/Database/PostgreSQL/LibPQ/FFI.hs b/src/Database/PostgreSQL/LibPQ/FFI.hs index 5bc11ac..0e87501 100644 --- a/src/Database/PostgreSQL/LibPQ/FFI.hs +++ b/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) diff --git a/test/Smoke.hs b/test/Smoke.hs index d21894b..8a8fae2 100644 --- a/test/Smoke.hs +++ b/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