From ea45394e55e5993fb1e5adbfcbd14f76226e765c Mon Sep 17 00:00:00 2001 From: VyacheslavHashov Date: Thu, 16 Feb 2017 21:22:55 +0300 Subject: [PATCH] Some refactoring --- src/Database/PostgreSQL/Driver/Error.hs | 73 +++++++++++-------- src/Database/PostgreSQL/Protocol/Decoders.hs | 37 +++------- .../PostgreSQL/Protocol/ExtractDataRows.hs | 5 +- src/Database/PostgreSQL/Protocol/Types.hs | 4 + 4 files changed, 60 insertions(+), 59 deletions(-) diff --git a/src/Database/PostgreSQL/Driver/Error.hs b/src/Database/PostgreSQL/Driver/Error.hs index 9ec5b14..505fff1 100644 --- a/src/Database/PostgreSQL/Driver/Error.hs +++ b/src/Database/PostgreSQL/Driver/Error.hs @@ -1,18 +1,51 @@ -module Database.PostgreSQL.Driver.Error where +module Database.PostgreSQL.Driver.Error + ( + -- * Errors + Error(..) + , AuthError(..) + -- * Exceptions + , ReceiverException(..) + , IncorrectUsage + , ProtocolException + -- * helpers + , throwIncorrectUsage + , throwProtocolEx + , eitherToProtocolEx + , throwErrorInIO + , throwAuthErrorInIO + ) where -import Control.Exception +import Control.Exception (throwIO, Exception(..), SomeException) import Data.ByteString (ByteString) -import qualified Data.ByteString.Char8 as BS import System.Socket (AddressInfoException) +import qualified Data.ByteString.Char8 as BS import Database.PostgreSQL.Protocol.Types (ErrorDesc) --- All possible exceptions: --- SocketException --- PeekException. --- ProtocolException --- IncorrectUsage. +-- All possible errors. +data Error + -- Error sended by PostgreSQL, not application error. + = PostgresError ErrorDesc + | AuthError AuthError + -- Receiver errors that may occur in receiver thread. + -- When such error occurs it means that receiver thread died. + | ReceiverError ReceiverException + deriving (Show) +-- | Unexcepted exception in the ReceiverThread. +newtype ReceiverException = ReceiverException SomeException + deriving (Show) + +-- Errors that might occur at authorization phase. +-- Non-recoverable. +data AuthError + = AuthNotSupported ByteString + | AuthInvalidAddress + | AuthAddressException AddressInfoException + deriving (Show) + +-- | Exception throwing when `readNextData` or `waitReadyForQuery` +-- used incorrectly. newtype IncorrectUsage = IncorrectUsage ByteString deriving (Show) @@ -20,6 +53,7 @@ instance Exception IncorrectUsage where displayException (IncorrectUsage msg) = "Incorrect usage: " ++ BS.unpack msg +-- | Exception in high-level parsing protocol messages. newtype ProtocolException = ProtocolException ByteString deriving (Show) @@ -36,29 +70,6 @@ throwProtocolEx = throwIO . ProtocolException eitherToProtocolEx :: Either ByteString a -> IO a eitherToProtocolEx = either throwProtocolEx pure --- All possible errors. -data Error - -- Error sended by PostgreSQL, not application error. - = PostgresError ErrorDesc - | AuthError AuthError - -- Receiver errors that may occur in receiver thread. When such error occur - -- it means that receiver thread died. - | ReceiverError ReceiverException - deriving (Show) - -newtype ReceiverException = ReceiverException SomeException - deriving (Show) - --- Errors that might occur at authorization phase. --- Non-recoverable. -data AuthError - = AuthNotSupported ByteString - | AuthInvalidAddress - | AuthAddressException AddressInfoException - deriving (Show) - --- Helpers - throwErrorInIO :: Error -> IO (Either Error a) throwErrorInIO = pure . Left diff --git a/src/Database/PostgreSQL/Protocol/Decoders.hs b/src/Database/PostgreSQL/Protocol/Decoders.hs index 5bfbba0..b662ce1 100644 --- a/src/Database/PostgreSQL/Protocol/Decoders.hs +++ b/src/Database/PostgreSQL/Protocol/Decoders.hs @@ -3,7 +3,7 @@ module Database.PostgreSQL.Protocol.Decoders ( -- * High-lever decoder - decodeNextServerMessage + decodeNextServerMessage -- * Decoders , decodeAuthResponse , decodeHeader @@ -14,26 +14,18 @@ module Database.PostgreSQL.Protocol.Decoders , parseErrorDesc ) where -import Control.Applicative -import Control.Monad import Data.Monoid ((<>)) import Data.Maybe (fromMaybe) import Data.Char (chr) -import Data.Word -import Foreign +import Data.Word (Word8, Word16, Word32) import Text.Read (readMaybe) import qualified Data.Vector as V import qualified Data.ByteString as B -import qualified Data.ByteString.Unsafe as B -import qualified Data.ByteString.Lazy as BL -import qualified Data.ByteString.Lazy.Internal as BL import Data.ByteString.Char8 as BS(readInteger, readInt, unpack, pack) import qualified Data.HashMap.Strict as HM import Database.PostgreSQL.Protocol.Types import Database.PostgreSQL.Protocol.Store.Decode -import Database.PostgreSQL.Protocol.Utils - -- | Parses and dispatches all server messages except `DataRow`. decodeNextServerMessage @@ -46,7 +38,7 @@ decodeNextServerMessage bs readMoreAction = go Nothing bs where -- Parse header go Nothing bs - | B.length bs < 5 = readMoreAndGo Nothing bs + | B.length bs < headerSize = readMoreAndGo Nothing bs | otherwise = let (rest, h) = runDecode decodeHeader bs in go (Just h) rest -- Parse body @@ -57,13 +49,15 @@ decodeNextServerMessage bs readMoreAction = go Nothing bs {-# INLINE readMoreAndGo #-} readMoreAndGo h = (go h =<<) . readMoreAction +-------------------------------- +-- Protocol decoders + decodeAuthResponse :: Decode AuthResponse decodeAuthResponse = do - c <- getWord8 - len <- getInt32BE + Header c len <- decodeHeader case chr $ fromIntegral c of 'E' -> AuthErrorResponse <$> - (getByteString (fromIntegral $ len - 4) >>= + (getByteString len >>= eitherToDecode .parseErrorDesc) 'R' -> do rType <- getInt32BE @@ -73,8 +67,7 @@ decodeAuthResponse = do 5 -> AuthenticationMD5Password . MD5Salt <$> getByteString 4 7 -> pure AuthenticationGSS 9 -> pure AuthenticationSSPI - 8 -> AuthenticationGSSContinue <$> - getByteString (fromIntegral $ len -8) + 8 -> AuthenticationGSSContinue <$> getByteString (len - 4) _ -> fail "Unknown authentication response" _ -> fail "Invalid auth response" @@ -115,14 +108,6 @@ decodeServerMessage (Header c len) = case chr $ fromIntegral c of rowsCount <- fromIntegral <$> getInt16BE RowDescription <$> V.replicateM rowsCount decodeFieldDescription --- | Decodes a single data value. Length `-1` indicates a NULL column value. --- No value bytes follow in the NULL case. -decodeValue :: Decode (Maybe B.ByteString) -decodeValue = getInt32BE >>= \n -> - if n == -1 - then pure Nothing - else Just <$> getByteString (fromIntegral n) - decodeTransactionStatus :: Decode TransactionStatus decodeTransactionStatus = getWord8 >>= \t -> case chr $ fromIntegral t of @@ -154,7 +139,8 @@ decodeFormat = getInt16BE >>= \f -> 1 -> pure Binary _ -> fail "Unknown field format" --- Parser that just work with B.ByteString, not Decode type +----------------------------- +-- Helper parsers that work with B.ByteString, not Decode type -- Helper to parse, not used by decoder itself parseServerVersion :: B.ByteString -> Either B.ByteString ServerVersion @@ -288,6 +274,7 @@ parseNoticeDesc s = do "is not presented in NoticeResponse message") Right . HM.lookup c +-- | Helper to lift Either in Decode eitherToDecode :: Either B.ByteString a -> Decode a eitherToDecode = either (fail . BS.unpack) pure diff --git a/src/Database/PostgreSQL/Protocol/ExtractDataRows.hs b/src/Database/PostgreSQL/Protocol/ExtractDataRows.hs index 0308abd..02f3638 100644 --- a/src/Database/PostgreSQL/Protocol/ExtractDataRows.hs +++ b/src/Database/PostgreSQL/Protocol/ExtractDataRows.hs @@ -27,8 +27,7 @@ loopExtractDataRows readMoreAction callback = go "" "" where go :: B.ByteString -> BL.ByteString -> IO () go bs acc - -- 5 - header size, defined by PostgreSQL - | B.length bs < 5 = readMoreAndGo bs acc + | B.length bs < headerSize = readMoreAndGo bs acc | otherwise = do ScanRowResult ch rest r <- scanDataRows bs -- We should force accumulator @@ -45,7 +44,7 @@ loopExtractDataRows readMoreAction callback = go "" "" -- that there are enough bytes to read header. 2 -> do Header mt len <- parseHeader rest - dispatchHeader mt len (B.drop 5 rest) newAcc + dispatchHeader mt len (B.drop headerSize rest) newAcc {-# INLINE dispatchHeader #-} dispatchHeader :: Word8 -> Int -> B.ByteString -> BL.ByteString -> IO () diff --git a/src/Database/PostgreSQL/Protocol/Types.hs b/src/Database/PostgreSQL/Protocol/Types.hs index 0e40df6..6da664c 100644 --- a/src/Database/PostgreSQL/Protocol/Types.hs +++ b/src/Database/PostgreSQL/Protocol/Types.hs @@ -154,6 +154,10 @@ data CancelRequest = CancelRequest !ServerProcessId !ServerSecretKey data Header = Header {-# UNPACK #-} !Word8 {-# UNPACK #-} !Int deriving (Show) +-- | Server message's header size. +headerSize :: Int +headerSize = 5 + -- | All possible responses from a server in usual query phase. data ServerMessage = BackendKeyData !ServerProcessId !ServerSecretKey