diff --git a/hpgsql.cabal b/hpgsql.cabal --- a/hpgsql.cabal +++ b/hpgsql.cabal @@ -1,7 +1,7 @@ cabal-version: 2.0 name: hpgsql -version: 0.2.0.0 +version: 0.2.0.1 synopsis: Pure Haskell PostgreSQL driver (no libpq) description: hpgsql is a pure Haskell implementation of a PostgreSQL driver category: Database diff --git a/src/Hpgsql/Internal.hs b/src/Hpgsql/Internal.hs --- a/src/Hpgsql/Internal.hs +++ b/src/Hpgsql/Internal.hs @@ -115,10 +115,12 @@ import Control.Exception.Safe (Exception (..), MonadThrow, SomeException, bracket, bracketOnError, finally, handleJust, mask, mask_, onException, throw, toException, tryJust) import Control.Monad (forM, forM_, join, unless, void, when) import Data.ByteString (ByteString) +import qualified Data.ByteString as BS import Data.ByteString.Internal (w2c) import qualified Data.ByteString.Lazy as LBS import Data.Data (Proxy (..)) import Data.Either (isLeft, isRight) +import Data.IORef (atomicWriteIORef, newIORef, readIORef) import Data.Int (Int32, Int64) import qualified Data.List as List import Data.List.NonEmpty (NonEmpty (..)) @@ -152,7 +154,6 @@ import qualified Network.Socket.ByteString.Lazy as SocketLBS import Streaming (Of (..), Stream) import qualified Streaming as S -import qualified Streaming.Internal as SInternal import qualified Streaming.Prelude as S import System.IO (stderr) import System.IO.Error (isResourceVanishedError) @@ -245,7 +246,7 @@ False -> throwIrrecoverableError "Socket is not marked as non-blocking, which is not supported by hpgsql. You might be running on an unsupported platform" True -> pure () socketIsClosed <- newMVar False - recvBuffer <- newMVar mempty + recvBuffer <- newIORef mempty sendBuffer <- newMVar mempty socketMutex <- mkMutex encodingContext <- newMVar (EncodingContext builtinPgTypesMap) @@ -261,7 +262,7 @@ preparedStatementNames = mempty, transactionStatusBeforeCurrentPipeline = TransIdle } - let hpgConnPartialDoNotReturn = HPgConnection sock socketIsClosed recvBuffer sendBuffer socketMutex originalConnStr addrInfo encodingContext connParams currentConnectionState 0 0 connOpts + let hpgConnPartialDoNotReturn = HPgConnection sock socketIsClosed recvBuffer sendBuffer socketMutex originalConnStr addrInfo encodingContext connParams currentConnectionState 0 "" connOpts case connectOrCancel of CancelNotConnect cancelRequest _ -> do nonAtomicSendMsg hpgConnPartialDoNotReturn cancelRequest @@ -366,7 +367,7 @@ addrCanonName = Nothing } else do - addrInfos <- Socket.getAddrInfo (Just Socket.defaultHints) (Just $ Text.unpack hostname) (Just $ show port) + addrInfos <- Socket.getAddrInfo (Just Socket.defaultHints {Socket.addrSocketType = Socket.Stream}) (Just $ Text.unpack hostname) (Just $ show port) case addrInfos of [] -> throwIrrecoverableError "Could not resolve address" addrInfo : _ -> pure addrInfo @@ -498,21 +499,31 @@ -- | Just like `receiveNextMsgWithMaskedContinuation` but passes a `Maybe a` to -- the continuation instead of throwing an exception on parser failure. On parsing -- failure, this makes sure the message buffer remains unaltered. -receiveNextMsgWithMaskedContinuation :: (Show a) => HPgConnection -> PgMsgParser a -> (a -> IO b) -> IO b +receiveNextMsgWithMaskedContinuation :: (Show a) => HPgConnection -> PgMsgParser a -> (a -> STM b) -> IO b receiveNextMsgWithMaskedContinuation conn parser f = receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure conn parser $ \case Right p -> f p Left (msgIdentChar, mPgError) -> throw IrrecoverableHpgsqlError {hpgsqlDetails = "Could not parse postgres message with ident char " <> Text.pack (show msgIdentChar) <> ". This is an internal error in Hpgsql. Please report it.", innerException = toException <$> mPgError, relatedStatement = Nothing} +data ReceiveWhat a b where + ReceiveDataRows :: ReceiveWhat DataRow [DataRow] + ReceiveArbitraryMsg :: PgMsgParser a -> (Either (Char, Maybe PostgresError) a -> STM b) -> ReceiveWhat a b + -- | Masks asynchronous exceptions in between the moment the message is extracted from -- the internal buffer and the supplied function runs to completion. -- This is important for control messages (i.e. not `DataRow`) because if you extract a -- message from the buffer, you must be able to update the connection's internal state, -- lest it will be left in a very broken place. --- CAREFUL: avoid doing networking or too much work in your supplied function. It must --- be really cheap! -receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure :: (Show a) => HPgConnection -> PgMsgParser a -> (Either (Char, Maybe PostgresError) a -> IO b) -> IO b -receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure conn@HPgConnection {socket, recvBuffer} parser f = do +receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure :: (Show a) => HPgConnection -> PgMsgParser a -> (Either (Char, Maybe PostgresError) a -> STM b) -> IO b +receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure conn parser f = receiveNextMsgGeneric conn (ReceiveArbitraryMsg parser f) + +-- | Masks asynchronous exceptions in between the moment the message is extracted from +-- the internal buffer and the supplied function runs to completion. +-- This is important for control messages (i.e. not `DataRow`) because if you extract a +-- message from the buffer, you must be able to update the connection's internal state, +-- lest it will be left in a very broken place. +receiveNextMsgGeneric :: (Show a) => HPgConnection -> ReceiveWhat a b -> IO b +receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do -- We need to preserve the invariant that the internal buffer's first byte is -- always the first byte of a valid Message while keeping this function -- interruptible. @@ -520,78 +531,99 @@ -- due to the presence of asynchronous exceptions. -- So we append to the buffer up until it has been fully fetched, -- and then extract it from the buffer in one piece. - currentBuf <- receiveUntilBufferHasAtLeast 5 - let bufLen = LBS.length currentBuf - let charAndLength = LBS.take 5 currentBuf + (initialBuf, initialBufLen) <- receiveUntilBufferHasAtLeast 5 + let charAndLength = LBS.take 5 initialBuf let (w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ LBS.uncons charAndLength lenLeftToFetch :: Int64 = fromIntegral $ either error id (Cereal.decodeLazy @Int32 lenbs) - 4 fullMessageLen = 5 + lenLeftToFetch - restOfMsg <- LBS.drop 5 . LBS.take fullMessageLen <$> if bufLen >= fullMessageLen then pure currentBuf else receiveUntilBufferHasAtLeast fullMessageLen - receivedNoticeOrParameterSoTryAgain <- go msgIdentChar restOfMsg fullMessageLen + (nowBuf, _nowBufLen) <- if initialBufLen >= fullMessageLen then pure (initialBuf, initialBufLen) else receiveUntilBufferHasAtLeast fullMessageLen + let restOfMsg = LBS.drop 5 $ LBS.take fullMessageLen nowBuf + receivedNoticeOrParameterSoTryAgain <- go msgIdentChar restOfMsg fullMessageLen nowBuf case receivedNoticeOrParameterSoTryAgain of - Nothing -> receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure conn parser f + Nothing -> receiveNextMsgGeneric conn receiveWhat Just res -> pure res where - -- We mask_ because if the supplied IO action runs with a message extracted from + modifyIORefIO ioref io = do + (!newC, !v) <- io + atomicWriteIORef ioref newC + pure v + -- We mask_ because if the supplied STM action runs with a message extracted from -- the recvBuffer, then we _must_ remove that message from recvBuffer. - -- TODO: These IO actions are always STM transactions (some times just a `pure` call), - -- so maybe we should reflect that in the type to avoid the impression that - -- anything can happen. - go msgIdentChar restOfMsg fullMessageLen = mask_ $ modifyMVar recvBuffer $ \lbs -> do - let !bufferWithoutMsg = - if LBS.length lbs >= fullMessageLen - then LBS.drop fullMessageLen lbs - else - error "Bug in Hpgsql. Internal buffer's bytes weren't filled enough" + -- Ideally we'd have non-retriable STM at the type-level here. Maybe later. + -- Make sure to do very little work inside `go`! + go msgIdentChar restOfMsg fullMessageLen nowBuf = mask_ $ modifyIORefIO recvBuffer $ do + let bufferWithoutMsg = LBS.drop fullMessageLen nowBuf + handleUnexpectedMsg onNotAnyReasonableMsg = + -- This could be a Notification, NOTICE or a ParameterStatus message, since these + -- can be received _at any time_ according to the docs. + case parsePgMessage msgIdentChar restOfMsg (Left3 <$> msgParser @NotificationResponse <|> Middle3 <$> msgParser @NoticeResponse <|> Right3 <$> msgParser @ParameterStatus) of + Just (Left3 notifResponse) -> do + debugPrint "Received notification. Will add it to internal queue." + STM.atomically $ do + sttv <- STM.readTVar $ internalConnectionState conn + STM.writeTQueue (notificationsReceived sttv) notifResponse + pure (bufferWithoutMsg, Nothing) + Just (Middle3 (NoticeResponse details)) -> do + let severity = + fromMaybe + "Notice of unknown severity" + (Map.lookup ErrorSeverity details) + humanmsg = fromMaybe "no message" (Map.lookup ErrorHumanReadableMsg details) + Text.hPutStrLn stderr $ decodeUtf8 $ LBS.toStrict $ severity <> ": " <> humanmsg + pure (bufferWithoutMsg, Nothing) + Just (Right3 (ParameterStatus {..})) -> do + when (parameterName == "client_encoding" && parameterValue /= "UTF8") $ throwIrrecoverableError $ "Postgres sent us a change of client_encoding to not UTF8, and Hpgsql only supports UTF8. The encoding postgres sent us is " <> parameterValue + modifyMVar_ (parameterStatusMap conn) $ \(!paramMap) -> pure (Map.insert parameterName parameterValue paramMap) + pure (bufferWithoutMsg, Nothing) + Nothing -> do + -- Just in case this is a postgres error, it might include useful information, + -- so we spit that out + let mPgError = mkPostgresError "" <$> parsePgMessage msgIdentChar restOfMsg (msgParser @ErrorResponse) + fmap (nowBuf,) $ Just <$> STM.atomically (onNotAnyReasonableMsg (msgIdentChar, mPgError)) + case receiveWhat of + ReceiveDataRows -> + -- Parse as many DataRows as we can to do as much work as we can per buffer "churn" + case Parser.parseOnly (Parser.matchLeftUnconsumed (Parser.parseMany customDataRowParser)) (LBS.toStrict nowBuf) of + Parser.ParseOk (unconsumedBuffer, msgs@(_ : _)) -> do + debugPrint $ "Received " ++ show msgs + pure (LBS.fromStrict unconsumedBuffer, Just msgs) + _ -> handleUnexpectedMsg $ const $ pure [] -- No error when we stop receiving DataRows, only emptiness + ReceiveArbitraryMsg parser f -> + case parsePgMessage msgIdentChar restOfMsg parser of + Just msg -> do + debugPrint $ "Received " ++ show msg + fmap (bufferWithoutMsg,) $ Just <$> STM.atomically (f (Right msg)) + Nothing -> handleUnexpectedMsg (f . Left) - case parsePgMessage msgIdentChar restOfMsg parser of - Just msg -> do - debugPrint $ "Received " ++ show msg - fmap (bufferWithoutMsg,) $ Just <$> f (Right msg) - Nothing -> do - -- This could be a Notification, NOTICE or a ParameterStatus message, since these - -- can be received _at any time_ according to the docs. - case parsePgMessage msgIdentChar restOfMsg (Left3 <$> msgParser @NotificationResponse <|> Middle3 <$> msgParser @NoticeResponse <|> Right3 <$> msgParser @ParameterStatus) of - Just (Left3 notifResponse) -> do - debugPrint "Received notification. Will add it to internal queue." - STM.atomically $ do - sttv <- STM.readTVar $ internalConnectionState conn - STM.writeTQueue (notificationsReceived sttv) notifResponse - pure (bufferWithoutMsg, Nothing) - Just (Middle3 (NoticeResponse details)) -> do - let severity = - fromMaybe - "Notice of unknown severity" - (Map.lookup ErrorSeverity details) - humanmsg = fromMaybe "no message" (Map.lookup ErrorHumanReadableMsg details) - Text.hPutStrLn stderr $ decodeUtf8 $ LBS.toStrict $ severity <> ": " <> humanmsg - pure (bufferWithoutMsg, Nothing) - Just (Right3 (ParameterStatus {..})) -> do - when (parameterName == "client_encoding" && parameterValue /= "UTF8") $ throwIrrecoverableError $ "Postgres sent us a change of client_encoding to not UTF8, and Hpgsql only supports UTF8. The encoding postgres sent us is " <> parameterValue - modifyMVar_ (parameterStatusMap conn) $ \(!paramMap) -> pure (Map.insert parameterName parameterValue paramMap) - pure (bufferWithoutMsg, Nothing) - Nothing -> - -- Just in case this is a postgres error, it might include useful information, - -- so we spit that out - let mPgError = mkPostgresError "" <$> parsePgMessage msgIdentChar restOfMsg (msgParser @ErrorResponse) - in fmap (lbs,) $ Just <$> f (Left (msgIdentChar, mPgError)) + -- Sadly we have to repeat the parsing of a DataRow message here, when it already + -- exists in the FromPgMessage instance and in the body of this function. Maybe + -- we can improve this later. + customDataRowParser = do + charAndLength <- Parser.take 5 + let (w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ BS.uncons charAndLength + lenLeftToFetch :: Int = fromIntegral $ either error id (Cereal.decode @Int32 lenbs) - 4 + if msgIdentChar == 'D' + then do + rowColumnData <- BS.drop 2 <$> Parser.take lenLeftToFetch + pure $ DataRow rowColumnData + else fail "Not a DataRow" -- \| Appends into the internal buffer by reading from the socket -- until the buffer has at least N bytes. - -- Returns the current buffer. - receiveUntilBufferHasAtLeast :: Int64 -> IO LBS.ByteString + -- Returns the current buffer and its length. + receiveUntilBufferHasAtLeast :: Int64 -> IO (LBS.ByteString, Int64) receiveUntilBufferHasAtLeast minBytesNecessary = do - currentBuffer <- readMVar recvBuffer + currentBuffer <- readIORef recvBuffer let nBytesInBuffer = LBS.length currentBuffer if nBytesInBuffer >= minBytesNecessary - then pure currentBuffer + then pure (currentBuffer, nBytesInBuffer) else do -- This takes from the kernel's recv buffer and appends to our buffer atomically, -- or an exception is thrown when receiving. - mask $ \restore -> rethrowAsIrrecoverable $ modifyMVar_ recvBuffer $ \lbs -> do + mask $ \restore -> rethrowAsIrrecoverable $ do restore $ socketWaitRead socket someBytes <- timeDebugNonBlockingOperation "recv" $ recvNonBlocking socket (max 16000 $ fromIntegral $ minBytesNecessary - nBytesInBuffer) - pure (lbs <> LBS.fromStrict someBytes) + atomicWriteIORef recvBuffer (currentBuffer <> LBS.fromStrict someBytes) receiveUntilBufferHasAtLeast minBytesNecessary sendCancellationRequest :: HPgConnection -> IO () @@ -618,7 +650,7 @@ () -- Already cancelled, no need to send another Nothing -> internalConnectOrCancel - (CancelNotConnect (CancelRequest (connPid conn) (cancelSecretKey conn)) (connectedTo conn)) + (CancelNotConnect (CancelRequest conn.connPid conn.cancelSecretKey) conn.connectedTo) (connOpts conn) (originalConnStr conn) (secondsToDiffTime 30) @@ -725,26 +757,26 @@ -- We do have an assertion in `updateQueryStateIfFirstOrThrow` to ensure we aren't -- dismissive of messages arriving in unexpected order. receiveParseOrBindCompleteAtomically = receiveNextMsgWithMaskedContinuation conn (Right3 <$> msgParser @ParseComplete <|> Left3 <$> msgParser @ErrorResponse <|> Middle3 <$> msgParser @BindComplete) $ \msgE -> do - STM.atomically $ updateQueryStateIfFirstOrThrow $ case msgE of + updateQueryStateIfFirstOrThrow $ case msgE of Right3 msg -> RespParseComplete msg Middle3 msg -> RespBindComplete msg Left3 err -> RespErrorResponse err receiveBindCompleteAtomically = receiveNextMsgWithMaskedContinuation conn (Right <$> msgParser @BindComplete <|> Left <$> msgParser @ErrorResponse) $ \msgE -> do - STM.atomically $ updateQueryStateIfFirstOrThrow $ case msgE of + updateQueryStateIfFirstOrThrow $ case msgE of Right msg -> RespBindComplete msg Left err -> RespErrorResponse err receiveNoDataOrRowDescriptionOrCopyInResponseAtomically = receiveNextMsgWithMaskedContinuation conn (Right <$> (Right3 <$> msgParser @RowDescription <|> Left3 <$> msgParser @NoData <|> Middle3 <$> msgParser @CopyInResponse) <|> Left <$> msgParser @ErrorResponse) $ \msgE -> do - STM.atomically $ updateQueryStateIfFirstOrThrow $ case msgE of + updateQueryStateIfFirstOrThrow $ case msgE of Right (Left3 msg) -> RespNoData msg Right (Middle3 msg) -> RespCopyInResponse msg Right (Right3 msg) -> RespRowDescription msg Left err -> RespErrorResponse err receiveDataRowOrCommandCompleteAtomically = receiveNextMsgWithMaskedContinuation conn (Right3 <$> msgParser @DataRow <|> Middle3 <$> msgParser @CommandComplete <|> Left3 <$> msgParser @ErrorResponse) $ \msgE -> do - STM.atomically $ updateQueryStateIfFirstOrThrow $ case msgE of + updateQueryStateIfFirstOrThrow $ case msgE of Right3 msg -> RespDataRow msg Middle3 msg -> RespCommandComplete msg Left3 msg -> RespErrorResponse msg - receiveReadyForQueryAtomically = receiveNextMsgWithMaskedContinuation conn (msgParser @ReadyForQuery) $ \rq -> STM.atomically $ updateQueryStateIfFirstOrThrow $ RespReadyForQuery rq + receiveReadyForQueryAtomically = receiveNextMsgWithMaskedContinuation conn (msgParser @ReadyForQuery) $ \rq -> updateQueryStateIfFirstOrThrow $ RespReadyForQuery rq getQueryStateIfFirstOrThrow :: STM QueryState getQueryStateIfFirstOrThrow = updateConnStateTxn conn $ \sttv -> do st <- STM.readTVar sttv @@ -865,33 +897,29 @@ pure (mERowDesc, pure $ Right cmd) (mERowDesc, Middle3 mDataRow) -> do let allOtherRows = - S.unfold - ( \() -> do - -- This is a bit ugly, but we try very carefully not to bear the cost of STM - -- transactions at all when receiving DataRows, since they can come in very large - -- amounts, but we need to make sure the _other_ important messages, like ErrorResponse - -- and CommandComplete, do update internal state. - mRow <- receiveNextMsgWithMaskedContinuationButDontThrowOnParsingFailure conn (msgParser @DataRow) pure - case mRow of - Right row -> pure $ Right (row :> ()) - Left _ -> do - stateAfterNextMsg <- snd <$> receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId - case stateAfterNextMsg of - ErrorResponseReceived _ err -> do - receiveReadyForQueryIfNecessary thisThreadId - pure $ Left $ Left err - CommandCompleteReceived _ cmd -> do - receiveReadyForQueryIfNecessary thisThreadId - pure $ Left $ Right cmd - ReadyForQueryReceived errOrCmd _ -> pure $ Left errOrCmd - _ -> throwIrrecoverableError "Internal error in Hpgsql. After a DataRow we should get either an ErrorResponse or a CommandComplete message" - ) - () + S.concat $ + S.unfold + ( \() -> do + mRow <- receiveNextMsgGeneric conn ReceiveDataRows + case mRow of + rows@(_ : _) -> pure $ Right (rows :> ()) + [] -> do + stateAfterNextMsg <- snd <$> receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId + case stateAfterNextMsg of + ErrorResponseReceived _ err -> do + receiveReadyForQueryIfNecessary thisThreadId + pure $ Left $ Left err + CommandCompleteReceived _ cmd -> do + receiveReadyForQueryIfNecessary thisThreadId + pure $ Left $ Right cmd + ReadyForQueryReceived errOrCmd _ -> pure $ Left errOrCmd + st -> throwIrrecoverableError $ "Internal error in Hpgsql. After the last DataRow we should get either an ErrorResponse or a CommandComplete message. State: " <> Text.pack (show st) + ) + () finalStream = case mDataRow of Nothing -> allOtherRows Just dr -> - SInternal.Step $ - dr :> allOtherRows + dr `S.cons` allOtherRows pure (mERowDesc, finalStream) where receiveReadyForQueryIfNecessary :: WeakThreadId -> IO () @@ -930,6 +958,9 @@ Right (CommandComplete n) -> pure n -- | Runs a query and streams results directly from the connection's socket, i.e. without using cursors. +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. -- -- Note on thread safety: it is important to note the same thread that runs this must -- be the thread that consumes the returned Stream, and the returned Stream must be @@ -939,6 +970,9 @@ queryS = querySWith rowDecoder -- | Runs a query and streams results directly from the connection's socket, i.e. without using cursors. +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. -- -- Note on thread safety: it is important to note the same thread that runs this must -- be the thread that consumes the returned Stream, and the returned Stream must be @@ -952,6 +986,10 @@ -- Prefer to use 'queryS' and 'querySWith', because 'RowDecoder' can typecheck PostgreSQL results -- even when no rows are returned by queries, and 'RowDecoderMonadic' cannot. -- +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. +-- -- Note on thread safety: it is important to note the same thread that runs this must -- be the thread that consumes the returned Stream, and the returned Stream must be -- consumed completely (up to the last row or a postgres error) before you are able @@ -1202,10 +1240,16 @@ Just tid -> (`elem` [ThreadDied, ThreadFinished]) <$> threadStatus tid -- | Fetches any number of rows by streaming them directly from the socket. +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. pipelineS :: (FromPgRow a) => Query -> Pipeline (IO (Stream (Of a) IO ())) pipelineS = pipelineSWith rowDecoder -- | Fetches any number of rows by streaming them directly from the socket, with a custom row decoder. +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. pipelineSWith :: RowDecoder a -> Query -> Pipeline (IO (Stream (Of a) IO ())) pipelineSWith rowparser@(RowDecoder _ _ expectedColFmts) (lastAndInitNE . breakQueryIntoStatements -> (firstQueriesToSend, lastQueryToSend)) = Pipeline @@ -1219,6 +1263,9 @@ -- | Prefer to use 'pipelineS' and 'pipelineSWith', because 'RowDecoder' can typecheck PostgreSQL results -- even when no rows are returned by queries, and 'RowDecoderMonadic' cannot. +-- Once you assign the returned Stream to a binding, you must consume that binding linearly, i.e. +-- you should never consume it more than once, for it _will_ give you wrong results if you do that. +-- This is normal behaviour for IO Streams, but it's worth emphasizing. pipelineSMWith :: RowDecoderMonadic a -> Query -> Pipeline (IO (Stream (Of a) IO ())) pipelineSMWith rowparser (lastAndInitNE . breakQueryIntoStatements -> (firstQueriesToSend, lastQueryToSend)) = Pipeline diff --git a/src/Hpgsql/InternalTypes.hs b/src/Hpgsql/InternalTypes.hs --- a/src/Hpgsql/InternalTypes.hs +++ b/src/Hpgsql/InternalTypes.hs @@ -72,6 +72,7 @@ #endif import qualified Control.Concurrent.STM as STM import Data.Hashable (hash) +import Data.IORef (IORef) import Data.Set (Set) import Hpgsql.Base (lastTwoAndInit, maximumOnOrDef, minimumOnOrDef) import Hpgsql.Builder (BinaryField) @@ -456,7 +457,7 @@ data HPgConnection = HPgConnection { socket :: !Socket, socketClosed :: !(MVar Bool), - recvBuffer :: !(MVar LBS.ByteString), + recvBuffer :: !(IORef LBS.ByteString), sendBuffer :: !(MVar [(LBS.ByteString, STM ())]), socketMutex :: !Mutex, originalConnStr :: !ConnectionString, @@ -465,7 +466,7 @@ parameterStatusMap :: !(MVar (Map Text Text)), internalConnectionState :: !(TVar InternalConnectionState), connPid :: !Int32, - cancelSecretKey :: !Int32, + cancelSecretKey :: !ByteString, connOpts :: !ConnectOpts } diff --git a/src/Hpgsql/Msgs.hs b/src/Hpgsql/Msgs.hs --- a/src/Hpgsql/Msgs.hs +++ b/src/Hpgsql/Msgs.hs @@ -87,14 +87,14 @@ PasswordMessage AuthenticationMethod String String deriving stock (Show) -data BackendKeyData = BackendKeyData {backendPid :: Int32, backendSecretKey :: Int32} +data BackendKeyData = BackendKeyData {backendPid :: !Int32, backendSecretKey :: !BS.ByteString} deriving stock (Show) data Bind = Bind {paramsValuesInOrder :: ![BinaryField], resultColumnFmts :: !Int, preparedStmtHash :: !(Maybe String)} deriving stock (Show) -- | PId first, secret key second -data CancelRequest = CancelRequest Int32 Int32 +data CancelRequest = CancelRequest Int32 ByteString deriving stock (Show) newtype CopyData = CopyData Builder @@ -160,9 +160,9 @@ _ -> Nothing instance FromPgMessage BackendKeyData where - msgParser = PgMsgParser $ \c (LBS.splitAt 4 -> (pidBS, secretBS)) -> case c of - 'K' -> case (,) <$> Cereal.decodeLazy @Int32 pidBS <*> Cereal.decodeLazy @Int32 secretBS of - Right (pid, secret) -> Just $ BackendKeyData {backendPid = pid, backendSecretKey = secret} + msgParser = PgMsgParser $ \c (LBS.splitAt 4 -> (pidBS, backendSecretKey)) -> case c of + 'K' -> case Cereal.decodeLazy @Int32 pidBS of + Right pid -> Just $ BackendKeyData {backendPid = pid, backendSecretKey = LBS.toStrict backendSecretKey} Left _ -> Nothing _ -> Nothing @@ -191,7 +191,7 @@ instance ToPgMessage CancelRequest where toPgMessage (CancelRequest pid secret) = - Builder.int32BE (4 + 4 + 4 + 4) <> Builder.int32BE 80877102 <> Builder.int32BE pid <> Builder.int32BE secret + Builder.int32BE (fromIntegral $ 4 + 4 + 4 + BS.length secret) <> Builder.int32BE 80877102 <> Builder.int32BE pid <> Builder.byteString secret instance ToPgMessage CopyData where toPgMessage (CopyData bs) = diff --git a/src/Hpgsql/SimpleParser.hs b/src/Hpgsql/SimpleParser.hs --- a/src/Hpgsql/SimpleParser.hs +++ b/src/Hpgsql/SimpleParser.hs @@ -12,6 +12,8 @@ take, endOfInput, match, + parseMany, + matchLeftUnconsumed, ) where @@ -32,7 +34,7 @@ (String -> r) -> -- \^ failure continuation (a -> ByteString -> r) -> - -- \^ success continuation + -- \^ success continuation, taking left-unparsed ByteString and parsed value r } @@ -85,6 +87,14 @@ ks mempty bs {-# INLINE take #-} +parseMany :: Parser a -> Parser [a] +parseMany p = Parser $ \bs' _kf ks -> let (vs, rest) = go bs' in ks vs rest + where + go bs = case parseOnly (matchLeftUnconsumed p) bs of + ParseOk (unconsumed, v) -> let (vs, rest) = go unconsumed in (v : vs, rest) + ParseFail _ -> ([], bs) +{-# INLINE parseMany #-} + -- | Succeeds only when the input has been fully consumed. endOfInput :: Parser () endOfInput = Parser $ \bs kf ks -> @@ -104,3 +114,14 @@ in ks (consumed, a) bs' ) {-# INLINE match #-} + +-- | Run a parser and additionally return the unconsumed/unparsed ByteString. +matchLeftUnconsumed :: Parser a -> Parser (ByteString, a) +matchLeftUnconsumed (Parser p) = Parser $ \bs kf ks -> + p + bs + kf + ( \a bs' -> + ks (bs', a) bs' + ) +{-# INLINE matchLeftUnconsumed #-}