mirror of
https://github.com/postgres-haskell/postgres-wire.git
synced 2024-11-22 05:53:12 +03:00
Improved Encode type
This commit is contained in:
VyacheslavHashov
parent
856d839a7c
commit
a0d12906c5
@ -5,6 +5,7 @@ module Database.PostgreSQL.Protocol.Encoders
|
||||
|
||||
import Data.Word (Word32)
|
||||
import Data.Monoid ((<>))
|
||||
import Data.Char (ord)
|
||||
import qualified Data.Vector as V
|
||||
import qualified Data.ByteString as B
|
||||
|
||||
@ -21,8 +22,11 @@ encodeStartMessage (StartupMessage (Username uname) (DatabaseName dbname))
|
||||
where
|
||||
len = fromIntegral $ getEncodeLen payload
|
||||
payload = putWord32BE currentVersion <>
|
||||
putPgString "user" <> putPgString uname <>
|
||||
putPgString "database" <> putPgString dbname <> putWord8 0
|
||||
putByteStringNull "user" <>
|
||||
putByteStringNull uname <>
|
||||
putByteStringNull "database" <>
|
||||
putByteStringNull dbname <>
|
||||
putWord8 0
|
||||
encodeStartMessage SSLRequest
|
||||
-- Value hardcoded by PostgreSQL docs.
|
||||
= putWord32BE 8 <> putWord32BE 80877103
|
||||
@ -31,8 +35,8 @@ encodeClientMessage :: ClientMessage -> Encode
|
||||
encodeClientMessage (Bind (PortalName portalName) (StatementName stmtName)
|
||||
paramFormat values resultFormat)
|
||||
= prependHeader 'B' $
|
||||
putPgString portalName <>
|
||||
putPgString stmtName <>
|
||||
putByteStringNull portalName <>
|
||||
putByteStringNull stmtName <>
|
||||
-- `1` means that the specified format code is applied to all parameters
|
||||
putWord16BE 1 <>
|
||||
encodeFormat paramFormat <>
|
||||
@ -43,32 +47,32 @@ encodeClientMessage (Bind (PortalName portalName) (StatementName stmtName)
|
||||
putWord16BE 1 <>
|
||||
encodeFormat resultFormat
|
||||
encodeClientMessage (CloseStatement (StatementName stmtName))
|
||||
= prependHeader 'C' $ putChar8 'S' <> putPgString stmtName
|
||||
= prependHeader 'C' $ putChar8 'S' <> putByteStringNull stmtName
|
||||
encodeClientMessage (ClosePortal (PortalName portalName))
|
||||
= prependHeader 'C' $ putChar8 'P' <> putPgString portalName
|
||||
= prependHeader 'C' $ putChar8 'P' <> putByteStringNull portalName
|
||||
encodeClientMessage (DescribeStatement (StatementName stmtName))
|
||||
= prependHeader 'D' $ putChar8 'S' <> putPgString stmtName
|
||||
= prependHeader 'D' $ putChar8 'S' <> putByteStringNull stmtName
|
||||
encodeClientMessage (DescribePortal (PortalName portalName))
|
||||
= prependHeader 'D' $ putChar8 'P' <> putPgString portalName
|
||||
= prependHeader 'D' $ putChar8 'P' <> putByteStringNull portalName
|
||||
encodeClientMessage (Execute (PortalName portalName) (RowsToReceive rows))
|
||||
= prependHeader 'E' $
|
||||
putPgString portalName <>
|
||||
putByteStringNull portalName <>
|
||||
putWord32BE rows
|
||||
encodeClientMessage Flush
|
||||
= prependHeader 'H' mempty
|
||||
encodeClientMessage (Parse (StatementName stmtName) (StatementSQL stmt) oids)
|
||||
= prependHeader 'P' $
|
||||
putPgString stmtName <>
|
||||
putPgString stmt <>
|
||||
putByteStringNull stmtName <>
|
||||
putByteStringNull stmt <>
|
||||
putWord16BE (fromIntegral $ V.length oids) <>
|
||||
foldMap (putWord32BE . unOid) oids
|
||||
encodeClientMessage (PasswordMessage passtext)
|
||||
= prependHeader 'p' $ putPgString $ getPassword passtext
|
||||
= prependHeader 'p' $ putByteStringNull $ getPassword passtext
|
||||
where
|
||||
getPassword (PasswordPlain p) = p
|
||||
getPassword (PasswordMD5 p) = p
|
||||
encodeClientMessage (SimpleQuery (StatementSQL stmt))
|
||||
= prependHeader 'Q' $ putPgString stmt
|
||||
= prependHeader 'Q' $ putByteStringNull stmt
|
||||
encodeClientMessage Sync
|
||||
= prependHeader 'S' mempty
|
||||
encodeClientMessage Terminate
|
||||
@ -76,18 +80,24 @@ encodeClientMessage Terminate
|
||||
|
||||
-- | Encodes single data values. Length `-1` indicates a NULL parameter value.
|
||||
-- No value bytes follow in the NULL case.
|
||||
{-# INLINE encodeValue #-}
|
||||
encodeValue :: Maybe B.ByteString -> Encode
|
||||
encodeValue Nothing = putWord32BE (-1)
|
||||
encodeValue (Just v) = putWord32BE (fromIntegral $ B.length v)
|
||||
<> putByteString v
|
||||
|
||||
{-# INLINE encodeFormat #-}
|
||||
encodeFormat :: Format -> Encode
|
||||
encodeFormat Text = putWord16BE 0
|
||||
encodeFormat Binary = putWord16BE 1
|
||||
|
||||
{-# INLINE prependHeader #-}
|
||||
prependHeader :: Char -> Encode -> Encode
|
||||
prependHeader c payload =
|
||||
-- Length includes itself but not the first message-type byte
|
||||
let len = 4 + fromIntegral (getEncodeLen payload)
|
||||
in putChar8 c <> putWord32BE len <> payload
|
||||
|
||||
{-# INLINE putChar8 #-}
|
||||
putChar8 :: Char -> Encode
|
||||
putChar8 = putWord8 . fromIntegral . ord
|
||||
|
||||
@ -3,10 +3,9 @@ module Database.PostgreSQL.Protocol.Store.Encode where
|
||||
import Data.Monoid (Monoid(..), (<>))
|
||||
import Foreign (poke, plusPtr, Ptr)
|
||||
import Data.Int (Int16, Int32)
|
||||
import Data.Word (Word8, Word16, Word32)
|
||||
import Data.Char (ord)
|
||||
import Data.Bits (shiftR)
|
||||
import Data.Word
|
||||
|
||||
import Foreign
|
||||
import Data.ByteString (ByteString)
|
||||
import Data.ByteString.Internal as B(toForeignPtr)
|
||||
import Data.Store.Core (Poke(..), unsafeEncodeWith, pokeStatePtr,
|
||||
@ -15,65 +14,78 @@ import Data.Store.Core (Poke(..), unsafeEncodeWith, pokeStatePtr,
|
||||
data Encode = Encode {-# UNPACK #-} !Int !(Poke ())
|
||||
|
||||
instance Monoid Encode where
|
||||
mempty = Encode 0 . Poke $ \_ offset -> pure (offset, ())
|
||||
{-# INLINE mempty #-}
|
||||
mempty = Encode 0 . Poke $ \_ offset -> pure (offset, ())
|
||||
|
||||
(Encode len1 f1) `mappend` (Encode len2 f2) = Encode (len1 + len2) (f1 *> f2)
|
||||
{-# INLINE mappend #-}
|
||||
(Encode len1 f1) `mappend` (Encode len2 f2) = Encode (len1 + len2) (f1 *> f2)
|
||||
|
||||
{-# INLINE getEncodeLen #-}
|
||||
getEncodeLen :: Encode -> Int
|
||||
getEncodeLen (Encode len _) = len
|
||||
{-# INLINE getEncodeLen #-}
|
||||
|
||||
{-# INLINE runEncode #-}
|
||||
runEncode :: Encode -> ByteString
|
||||
runEncode (Encode len f) = unsafeEncodeWith f len
|
||||
{-# INLINE runEncode #-}
|
||||
|
||||
fixedPrim :: Int -> (Ptr Word8 -> IO ()) -> Encode
|
||||
fixedPrim len f = Encode len . Poke $ \state offset -> do
|
||||
{-# INLINE fixed #-}
|
||||
fixed :: Int -> (Ptr Word8 -> IO ()) -> Encode
|
||||
fixed len f = Encode len . Poke $ \state offset -> do
|
||||
f $ pokeStatePtr state `plusPtr` offset
|
||||
let !newOffset = offset + len
|
||||
return (newOffset, ())
|
||||
{-# INLINE fixedPrim #-}
|
||||
|
||||
putWord8 :: Word8 -> Encode
|
||||
putWord8 w = fixedPrim 1 $ \p -> poke p w
|
||||
{-# INLINE putWord8 #-}
|
||||
|
||||
putChar8 :: Char -> Encode
|
||||
putChar8 = putWord8 . fromIntegral . ord
|
||||
{-# INLINE putChar8 #-}
|
||||
|
||||
putWord16BE :: Word16 -> Encode
|
||||
putWord16BE w = fixedPrim 2 $ \p -> do
|
||||
poke p (fromIntegral (shiftR w 8) :: Word8)
|
||||
poke (p `plusPtr` 1) (fromIntegral w :: Word8)
|
||||
{-# INLINE putWord16BE #-}
|
||||
|
||||
putWord32BE :: Word32 -> Encode
|
||||
putWord32BE w = fixedPrim 4 $ \p -> do
|
||||
poke p (fromIntegral (shiftR w 24) :: Word8)
|
||||
poke (p `plusPtr` 1) (fromIntegral (shiftR w 16) :: Word8)
|
||||
poke (p `plusPtr` 2) (fromIntegral (shiftR w 8) :: Word8)
|
||||
poke (p `plusPtr` 3) (fromIntegral w :: Word8)
|
||||
{-# INLINE putWord32BE #-}
|
||||
|
||||
putInt32BE :: Int32 -> Encode
|
||||
putInt32BE = putWord32BE . fromIntegral
|
||||
{-# INLINE putInt32BE #-}
|
||||
|
||||
putInt16BE :: Int16 -> Encode
|
||||
putInt16BE = putWord16BE . fromIntegral
|
||||
{-# INLINE putInt16BE #-}
|
||||
|
||||
{-# INLINE putByteString #-}
|
||||
putByteString :: ByteString -> Encode
|
||||
putByteString bs =
|
||||
let (ptr, offset, len) = toForeignPtr bs
|
||||
in Encode len $ pokeFromForeignPtr ptr offset len
|
||||
{-# INLINE putByteString #-}
|
||||
|
||||
-- | C-like string
|
||||
putPgString :: ByteString -> Encode
|
||||
putPgString bs = putByteString bs <> putWord8 0
|
||||
{-# INLINE putPgString #-}
|
||||
{-# INLINE putByteStringNull #-}
|
||||
putByteStringNull :: ByteString -> Encode
|
||||
putByteStringNull bs = putByteString bs <> putWord8 0
|
||||
|
||||
{-# INLINE putWord8 #-}
|
||||
putWord8 :: Word8 -> Encode
|
||||
putWord8 w = fixed 1 $ \p -> poke p w
|
||||
|
||||
{-# INLINE putWord16BE #-}
|
||||
putWord16BE :: Word16 -> Encode
|
||||
putWord16BE w = fixed 2 $ \p -> poke (castPtr p) (byteSwap16 w)
|
||||
|
||||
{-# INLINE putWord32BE #-}
|
||||
putWord32BE :: Word32 -> Encode
|
||||
putWord32BE w = fixed 4 $ \p -> poke (castPtr p) (byteSwap32 w)
|
||||
|
||||
{-# INLINE putWord64BE #-}
|
||||
putWord64BE :: Word64 -> Encode
|
||||
putWord64BE w = fixed 8 $ \p -> poke (castPtr p) (byteSwap64 w)
|
||||
|
||||
{-# INLINE putInt16BE #-}
|
||||
putInt16BE :: Int16 -> Encode
|
||||
putInt16BE = putWord16BE . fromIntegral
|
||||
|
||||
{-# INLINE putInt32BE #-}
|
||||
putInt32BE :: Int32 -> Encode
|
||||
putInt32BE = putWord32BE . fromIntegral
|
||||
|
||||
{-# INLINE putInt64BE #-}
|
||||
putInt64BE :: Int64 -> Encode
|
||||
putInt64BE = putWord64BE . fromIntegral
|
||||
|
||||
{-# INLINE putFloat32BE #-}
|
||||
putFloat32BE :: Float -> Encode
|
||||
putFloat32BE float = fixed 4 $ \ptr -> byteSwap32 <$> floatToWord float
|
||||
>>= poke (castPtr ptr)
|
||||
|
||||
{-# INLINE putFloat64BE #-}
|
||||
putFloat64BE :: Double -> Encode
|
||||
putFloat64BE double = fixed 8 $ \ptr -> byteSwap64 <$> floatToWord double
|
||||
>>= poke (castPtr ptr)
|
||||
|
||||
{-# INLINE floatToWord #-}
|
||||
floatToWord :: (Storable word, Storable float) => float -> IO word
|
||||
floatToWord float = alloca $ \buf -> do
|
||||
poke (castPtr buf) float
|
||||
peek buf
|
||||
|
||||
Loading…
Reference in New Issue
Block a user