diff --git a/Database/Monarch.hs b/Database/Monarch.hs --- a/Database/Monarch.hs +++ b/Database/Monarch.hs @@ -1,111 +1,532 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} --- | This module makes tokyotyrant-haskell monadic. +-- | This module provide TokyoTyrant monadic access interface. +-- +-- TokyoTyrant Original Binary Protocol() is implemented. module Database.Monarch ( - -- * Basic Types - Key, Value, Host, Port - -- * TokyoTyrant DB action monad and its runner - , Monarch, runMonarch - -- * Fetch/store single key/value - , get, put, putKeep - -- * Fetch/store multiple key/value - , mget, mput - -- * Delete key/value - , delete, vanish - -- * Search key prefix - , fwmkeys + Monarch, runMonarch + , ExtOption(..), RestoreOption(..), MiscOption(..) + , Code(..) + , put, putKeep, putCat, putShiftLeft + , putNoResponse + , out + , get, multipleGet + , valueSize + , iterInit, iterNext + , forwardMatchingKeys + , addInt, addDouble + , ext, sync, optimize, vanish, copy, restore + , setMaster + , recordNum, size + , status + , misc ) where -import qualified Database.TokyoTyrant.FFI as TT -import Data.ByteString +import Data.Bits +import Data.Int +import Data.Conduit +import Data.Conduit.Network +import qualified Data.Conduit.Binary as CB +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS import Control.Applicative import Control.Monad.Reader import Control.Monad.Error -import Control.Exception +import Data.IORef +import qualified Data.Binary as B +import Data.Binary.Put (runPut, putWord32be, putByteString) +import Data.Binary.Get (runGet, getWord32be) -type Key = ByteString +type TTPipe = + Pipe BS.ByteString BS.ByteString BS.ByteString () IO -type Value = ByteString +-- | Error code +data Code = Success + | InvalidOperation + | HostNotFound + | ConnectionRefused + | SendError + | ReceiveError + | ExistingRecord + | NoRecordFound + | MiscellaneousError + deriving (Eq, Show) -type Host = ByteString +instance Error Code -type Port = Int +class BitFlag32 a where + fromOption :: a -> Int32 +-- | Options for scripting extension +data ExtOption = RecordLocking -- ^ record locking + | GlobalLocking -- ^ global locking + +instance BitFlag32 ExtOption where + fromOption RecordLocking = 0x1 + fromOption GlobalLocking = 0x2 + +-- | Options for restore +data RestoreOption = ConsistencyChecking -- ^ consistency checking + +instance BitFlag32 RestoreOption where + fromOption ConsistencyChecking = 0x1 + +-- | Options for miscellaneous operation +data MiscOption = NoUpdateLog -- ^ omission of update log + +instance BitFlag32 MiscOption where + fromOption NoUpdateLog = 0x1 + +toCode :: Int -> Code +toCode 0 = Success +toCode 1 = InvalidOperation +toCode 2 = HostNotFound +toCode 3 = ConnectionRefused +toCode 4 = SendError +toCode 5 = ReceiveError +toCode 6 = ExistingRecord +toCode 7 = NoRecordFound +toCode 9999 = MiscellaneousError +toCode _ = error "Invalid Code" + +putMagic :: B.Word8 + -> B.Put +putMagic magic = B.putWord8 0xC8 >> B.putWord8 magic + +putOptions :: BitFlag32 option => + [option] + -> B.Put +putOptions = putWord32be . fromIntegral . + foldl (.|.) 0 . map fromOption + + -- | A monad supporting TokyoTyrant access. -newtype Monarch a = Monarch { unMonarch :: ErrorT String (ReaderT TT.Connection IO) a} +newtype Monarch a = + Monarch { unMonarch :: ErrorT Code TTPipe a } deriving ( Functor, Applicative, Monad, MonadIO - , MonadReader TT.Connection, MonadError String ) + , MonadError Code ) -- | Run Monarch with TokyoTyrant at target host and port. runMonarch :: MonadIO m => - Host -- ^ host - -> Port -- ^ port - -> Monarch a -- ^ action - -> m (Either String a) -runMonarch host port monarch = liftIO $ bracket open' close' action + String + -> Int + -> Monarch a + -> m (Either Code a) +runMonarch host port action = + liftIO $ do + result <- liftIO $ newIORef $ Left Success + let remote = ClientSettings port host + runTCPClient remote $ client action result + readIORef result + +client :: Monarch a + -> IORef (Either Code a) + -> Application IO +client action result src sink = src $$ conduit =$ sink where - action = (return . Left) `either` (runReaderT . runErrorT . unMonarch $ monarch) - open' = TT.open host port - close' = (const $ return ()) `either` TT.close + conduit = runErrorT (unMonarch action) >>= + liftIO . writeIORef result -liftMonarch :: (MonadIO m, MonadError a m) => IO (Either a b) -> m b -liftMonarch action = liftIO action >>= either throwError return +lengthBS32 :: BS.ByteString -> B.Word32 +lengthBS32 = fromIntegral . BS.length --- | Get a value. If not found, return Nothing. -get :: Key -- ^ key - -> Monarch (Maybe Value) -get key = do - conn <- ask - liftMonarch $ TT.get conn key +fromLBS :: LBS.ByteString -> BS.ByteString +fromLBS = BS.pack . LBS.unpack --- | Destructive store a value. -put :: Key -- ^ key - -> Value -- ^ value +liftMonarch :: TTPipe a + -> Monarch a +liftMonarch = Monarch . lift + +yieldRequest :: B.Put -> Monarch () +yieldRequest = liftMonarch . yield . fromLBS . runPut + +responseCode :: Monarch Code +responseCode = + liftMonarch CB.head >>= + maybe (throwError MiscellaneousError) + (return . toCode . fromIntegral) + +parseLBS :: Monarch LBS.ByteString +parseLBS = liftMonarch $ + CB.take 4 >>= + CB.take . fromIntegral . runGet getWord32be + +parseBS :: Monarch BS.ByteString +parseBS = fromLBS <$> parseLBS + +parseWord32 :: Monarch B.Word32 +parseWord32 = liftMonarch (CB.take 4) >>= + return . runGet getWord32be + +parseInt64 :: Monarch Int64 +parseInt64 = liftMonarch (CB.take 8) >>= + return . runGet (B.get :: B.Get Int64) + +parseDouble :: Monarch Double +parseDouble = do + integ <- fromIntegral <$> parseInt64 + fract <- fromIntegral <$> parseInt64 + return $ integ + fract * 1e-12 + +parseKeyValue :: Monarch (BS.ByteString, BS.ByteString) +parseKeyValue = + liftMonarch $ do + ksiz <- CB.take 4 + vsiz <- CB.take 4 + key <- CB.take . fromIntegral $ + runGet getWord32be ksiz + value <- CB.take . fromIntegral $ + runGet getWord32be vsiz + return (fromLBS key, fromLBS value) + +communicate :: B.Put + -> (Code -> Monarch a) + -> Monarch a +communicate makeRequest makeResponse = + yieldRequest makeRequest >> + responseCode >>= + makeResponse + +-- | Store a record. +-- If a record with the same key exists in the database, +-- it is overwritten. +put :: BS.ByteString -- ^ key + -> BS.ByteString -- ^ value -> Monarch () -put key value = do - conn <- ask - liftMonarch $ TT.put conn key value +put key value = communicate request response + where + request = do + putMagic 0x10 + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putByteString key + putByteString value + response Success = return () + response code = throwError code --- | Non-Destructive store a value. -putKeep :: Key -- ^ key - -> Value -- ^ value +-- | Store a new record. +-- If a record with the same key exists in the database, +-- this function has no effect. +putKeep :: BS.ByteString -- ^ key + -> BS.ByteString -- ^ value -> Monarch () -putKeep key value = do - conn <- ask - liftMonarch $ TT.putKeep conn key value +putKeep key value = communicate request response + where + request = do + putMagic 0x11 + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putByteString key + putByteString value + response Success = return () + response InvalidOperation = return () + response code = throwError code --- | Get values. -mget :: [Key] -- ^ keys to get - -> Monarch [(Key, Value)] -mget keys = do - conn <- ask - liftMonarch $ TT.mget conn keys +-- | Concatenate a value at the end of the existing record. +-- If there is no corresponding record, a new record is created. +putCat :: BS.ByteString -- ^ key + -> BS.ByteString -- ^ value + -> Monarch () +putCat key value = communicate request response + where + request = do + putMagic 0x12 + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putByteString key + putByteString value + response Success = return () + response code = throwError code --- | Put values. -mput :: [(Key, Value)] -- ^ key/value pairs to put - -> Monarch () -mput kvs = do - conn <- ask - liftMonarch $ TT.mput conn kvs +-- | Concatenate a value at the end of the existing record and shift it to the left. +-- If there is no corresponding record, a new record is created. +putShiftLeft :: BS.ByteString -- ^ key + -> BS.ByteString -- ^ value + -> Int -- ^ width + -> Monarch () +putShiftLeft key value width = communicate request response + where + request = do + putMagic 0x13 + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putWord32be $ fromIntegral width + putByteString key + putByteString value + response Success = return () + response code = throwError code --- | Delete a value. -delete :: Key -- ^ key to delete - -> Monarch () -delete key = do - conn <- ask - liftMonarch $ TT.delete conn key +-- | Store a record without response. +-- If a record with the same key exists in the database, it is overwritten. +putNoResponse :: BS.ByteString -- ^ key + -> BS.ByteString -- ^ value + -> Monarch () +putNoResponse key value = yieldRequest request + where + request = do + putMagic 0x18 + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putByteString key + putByteString value --- | Delete all value. +-- | Remove a record. +out :: BS.ByteString -- ^ key + -> Monarch () +out key = communicate request response + where + request = do + putMagic 0x20 + putWord32be $ lengthBS32 key + putByteString key + response Success = return () + response InvalidOperation = return () + response code = throwError code + +-- | Retrieve a record. +get :: BS.ByteString -- ^ key + -> Monarch (Maybe BS.ByteString) +get key = communicate request response + where + request = do + putMagic 0x30 + putWord32be $ lengthBS32 key + putByteString key + response Success = Just <$> parseBS + response InvalidOperation = return Nothing + response code = throwError code + +-- | Retrieve records. +multipleGet :: [BS.ByteString] -- ^ keys + -> Monarch [(BS.ByteString, BS.ByteString)] +multipleGet keys = communicate request response + where + request = do + putMagic 0x31 + putWord32be $ fromIntegral $ length keys + mapM_ (\key -> do + putWord32be $ lengthBS32 key + putByteString key) keys + response Success = do + siz <- fromIntegral <$> parseWord32 + sequence $ Prelude.replicate siz parseKeyValue + response code = throwError code + +-- | Get the size of the value of a record. +valueSize :: BS.ByteString -- ^ key + -> Monarch (Maybe Int) +valueSize key = communicate request response + where + request = do + putMagic 0x38 + putWord32be $ lengthBS32 key + putByteString key + response Success = Just . fromIntegral <$> parseWord32 + response InvalidOperation = return Nothing + response code = throwError code + +-- | Initialize the iterator. +iterInit :: Monarch () +iterInit = communicate request response + where + request = putMagic 0x50 + response Success = return () + response code = throwError code + +-- | Get the next key of the iterator. +-- The iterator can be updated by multiple connections and then it is not assured that every record is traversed. +iterNext :: Monarch (Maybe BS.ByteString) +iterNext = communicate request response + where + request = putMagic 0x51 + response Success = Just <$> parseBS + response InvalidOperation = return Nothing + response code = throwError code + +-- | Get forward matching keys. +forwardMatchingKeys :: BS.ByteString -- ^ key prefix + -> Int -- ^ maximum number of keys to be fetched + -> Monarch [BS.ByteString] +forwardMatchingKeys prefix n = communicate request response + where + request = do + putMagic 0x58 + putWord32be $ lengthBS32 prefix + putWord32be $ fromIntegral n + putByteString prefix + response Success = do + siz <- fromIntegral <$> parseWord32 + sequence $ Prelude.replicate siz parseBS + response code = throwError code + +-- | Add an integer to a record. +-- If the corresponding record exists, the value is treated as an integer and is added to. +-- If no record corresponds, a new record of the additional value is stored. +addInt :: BS.ByteString -- ^ key + -> Int -- ^ value + -> Monarch Int +addInt key n = communicate request response + where + request = do + putMagic 0x60 + putWord32be $ lengthBS32 key + putWord32be $ fromIntegral n + putByteString key + response Success = fromIntegral <$> parseWord32 + response code = throwError code + +-- | Add a real number to a record. +-- If the corresponding record exists, the value is treated as a real number and is added to. +-- If no record corresponds, a new record of the additional value is stored. +addDouble :: BS.ByteString -- ^ key + -> Double -- ^ value + -> Monarch Double +addDouble key n = communicate request response + where + request = do + putMagic 0x61 + putWord32be $ lengthBS32 key + B.put (truncate n :: Int64) + B.put (truncate (snd (properFraction n :: (Int,Double)) * 1e12) :: Int64) + putByteString key + response Success = parseDouble + response code = throwError code + +-- | Call a function of the script language extension. +ext :: BS.ByteString -- ^ function + -> [ExtOption] -- ^ option flags + -> BS.ByteString -- ^ key + -> BS.ByteString -- ^ value + -> Monarch BS.ByteString +ext func opts key value = communicate request response + where + request = do + putMagic 0x68 + putWord32be $ lengthBS32 func + putOptions opts + putWord32be $ lengthBS32 key + putWord32be $ lengthBS32 value + putByteString func + putByteString key + putByteString value + response Success = parseBS + response code = throwError code + +-- | Synchronize updated contents with the file and the device. +sync :: Monarch () +sync = communicate request response + where + request = putMagic 0x70 + response Success = return () + response code = throwError code + +-- | Optimize the storage. +optimize :: BS.ByteString -- ^ parameter + -> Monarch () +optimize param = communicate request response + where + request = do + putMagic 0x71 + putWord32be $ lengthBS32 param + putByteString param + response Success = return () + response code = throwError code + +-- | Remove all records. vanish :: Monarch () -vanish = do - conn <- ask - liftMonarch $ TT.vanish conn +vanish = communicate request response + where + request = putMagic 0x72 + response Success = return () + response code = throwError code --- | Search keys by prefix. -fwmkeys :: ByteString -- ^ key prefix - -> Maybe Int -- ^ max # of returned keys; - -- Nothing means no limit - -> Monarch [Key] -fwmkeys key n = do - conn <- ask - liftMonarch $ TT.fwmkeys conn key (maybe (-1) id n) +-- | Copy the database file. +copy :: BS.ByteString -- ^ path + -> Monarch () +copy path = communicate request response + where + request = do + putMagic 0x73 + putWord32be $ lengthBS32 path + putByteString path + response Success = return () + response code = throwError code + +-- | Restore the database file from the update log. +restore :: Integral a => + BS.ByteString -- ^ path + -> a -- ^ beginning time stamp in microseconds + -> [RestoreOption] -- ^ option flags + -> Monarch () +restore path usec opts = communicate request response + where + request = do + putMagic 0x74 + putWord32be $ lengthBS32 path + B.put (fromIntegral usec :: Int64) + putOptions opts + putByteString path + response Success = return () + response code = throwError code + +-- | Set the replication master. +setMaster :: Integral a => + BS.ByteString -- ^ host + -> Int -- ^ port + -> a -- ^ beginning time stamp in microseconds + -> [RestoreOption] -- ^ option flags + -> Monarch () +setMaster host port usec opts = communicate request response + where + request = do + putMagic 0x78 + putWord32be $ lengthBS32 host + putWord32be $ fromIntegral port + B.put (fromIntegral usec :: Int64) + putOptions opts + putByteString host + response Success = return () + response code = throwError code + +-- | Get the number of records. +recordNum :: Monarch Int64 +recordNum = communicate request response + where + request = putMagic 0x80 + response Success = parseInt64 + response code = throwError code + +-- | Get the size of the database. +size :: Monarch Int64 +size = communicate request response + where + request = putMagic 0x81 + response Success = parseInt64 + response code = throwError code + +-- | Get the status string of the database. +status :: Monarch BS.ByteString +status = communicate request response + where + request = putMagic 0x88 + response Success = parseBS + response code = throwError code + +-- | Call a versatile function for miscellaneous operations. +misc :: BS.ByteString -- ^ function name + -> [MiscOption] -- ^ option flags + -> [BS.ByteString] -- ^ arguments + -> Monarch [BS.ByteString] +misc func opts args = communicate request response + where + request = do + putMagic 0x90 + putWord32be $ lengthBS32 func + putOptions opts + putWord32be $ fromIntegral $ length args + putByteString func + mapM_ putByteString args + response Success = do + siz <- fromIntegral <$> parseWord32 + sequence $ Prelude.replicate siz parseBS + response code = throwError code diff --git a/monarch.cabal b/monarch.cabal --- a/monarch.cabal +++ b/monarch.cabal @@ -1,5 +1,5 @@ name: monarch -version: 0.1.0.0 +version: 0.2.0.0 synopsis: Monadic interface for TokyoTyrant. description: This package provides simple monadic interface for TokyoTyrant. license: BSD3 @@ -17,7 +17,9 @@ build-depends: base ==4.* , mtl ==2.1.* , bytestring ==0.9.* - , tokyotyrant-haskell ==1.0.* + , conduit ==0.5.* + , network-conduit ==0.5.* + , binary ==0.5.* test-suite unit-test hs-source-dirs: test, . @@ -26,7 +28,9 @@ build-depends: base ==4.* , mtl ==2.1.* , bytestring ==0.9.* - , tokyotyrant-haskell ==1.0.* + , conduit ==0.5.* + , network-conduit ==0.5.* + , binary ==0.5.* , hspec ==1.3.* , HUnit ==1.2.* diff --git a/test/test.hs b/test/test.hs --- a/test/test.hs +++ b/test/test.hs @@ -1,31 +1,243 @@ {-# LANGUAGE OverloadedStrings #-} import Database.Monarch +import Control.Applicative +import Data.List +import Data.ByteString.Char8 () import Test.HUnit import Test.Hspec.Monadic import Test.Hspec.HUnit () main :: IO () main = hspec $ do - describe "TokyoTyrant" $ do - it "get unstored value" caseGetUnstored - it "get already stored value" caseGetStored + describe "put" $ do + it "store a record" casePutRecord + it "overwrite a record if same key exists" casePutOverwriteRecord + describe "putkeep" $ do + it "store a new record" casePutKeepNewRecord + it "has no effect if same key exists" casePutKeepNoEffect + describe "putcat" $ do + it "concatenate a value at the end of the existing record" casePutCatRecord + it "store a new record if there is no corresponding record" casePutCatNewRecord + describe "putshl" $ do + it "concatenate a value at the end of the existing record and shift it to the left" casePutShlRecord + it "store a new record if there is no corresponding record" casePutShlNewRecord + describe "putnr" $ do + it "store a record" casePutNrRecord + it "overwrite a record if same key exists" casePutNrOverwriteRecord + describe "out" $ do + it "remove a record" caseOutRecord + it "no effect if same key not exists" caseOutNoEffect + describe "get" $ do + it "retrieve a record" caseGetRecord + describe "mget" $ do + it "retrieve records" caseMgetRecords + describe "vsiz" $ do + it "get the size of the value of a record" caseVsizRecord + describe "iterinit" $ do + it "initialize the iterator" caseIterinit + describe "iternext" $ do + it "get the next key of the iterator" caseIternext + it "invalid if end iterator" caseIternextInvalid + describe "fwmkeys" $ do + it "get forward matching keys" caseFwmkeys -caseGetUnstored :: Assertion -caseGetUnstored = do +returns :: (Eq a, Show a) => + Monarch a + -> Either Code a + -> IO () +action `returns` expected = do result <- runMonarch "localhost" 1978 $ do - vanish - result <- get "foo" - vanish - return result - result @?= Right Nothing + vanish + result <- action + vanish + return result + result @?= expected -caseGetStored :: Assertion -caseGetStored = do - result <- runMonarch "localhost" 1978 $ do - vanish - put "foo" "bar" - result <- get "foo" - vanish - return result - result @?= Right (Just "bar") +casePutRecord :: Assertion +casePutRecord = + action `returns` Right (Just "bar") + where + action = do + put "foo" "bar" + get "foo" + +casePutOverwriteRecord :: Assertion +casePutOverwriteRecord = + action `returns` Right (Just "hoge") + where + action = do + put "foo" "bar" + put "foo" "hoge" + get "foo" + +casePutKeepNewRecord :: Assertion +casePutKeepNewRecord = + action `returns` Right (Just "hoge") + where + action = do + putKeep "foo" "hoge" + get "foo" + +casePutKeepNoEffect :: Assertion +casePutKeepNoEffect = + action `returns` Right (Just "bar") + where + action = do + putKeep "foo" "bar" + putKeep "foo" "hoge" + get "foo" + +casePutCatRecord :: Assertion +casePutCatRecord = + action `returns` Right (Just "abracadabra") + where + action = do + put "foo" "abra" + putCat "foo" "cadabra" + get "foo" + +casePutCatNewRecord :: Assertion +casePutCatNewRecord = + action `returns` Right (Just "cadabra") + where + action = do + putCat "foo" "cadabra" + get "foo" + +casePutShlRecord :: Assertion +casePutShlRecord = + action `returns` Right (Just "racadabra") + where + action = do + put "foo" "abra" + putShiftLeft "foo" "cadabra" 9 + get "foo" + +casePutShlNewRecord :: Assertion +casePutShlNewRecord = + action `returns` Right (Just "cadabra") + where + action = do + putShiftLeft "foo" "cadabra" 4 + get "foo" + +casePutNrRecord :: Assertion +casePutNrRecord = + action `returns` Right (Just "bar") + where + action = do + putNoResponse "foo" "bar" + get "foo" + +casePutNrOverwriteRecord :: Assertion +casePutNrOverwriteRecord = + action `returns` Right (Just "hoge") + where + action = do + putNoResponse "foo" "bar" + putNoResponse "foo" "hoge" + get "foo" + +caseOutRecord :: Assertion +caseOutRecord = + action `returns` Right (Just "bar", Nothing) + where + action = do + put "foo" "bar" + put "hoge" "fuga" + out "hoge" + stored <- get "foo" + unstored <- get "hoge" + return (stored, unstored) + +caseOutNoEffect :: Assertion +caseOutNoEffect = + action `returns` Right () + where + action = do + out "hoge" + +caseGetRecord :: Assertion +caseGetRecord = + action `returns` Right (Just "bar", Nothing) + where + action = do + put "foo" "bar" + stored <- get "foo" + unstored <- get "bar" + return (stored, unstored) + +caseMgetRecords :: Assertion +caseMgetRecords = + action `returns` Right [ ("foo", "bar") + , ("huga", "hoge") + , ("abra", "cadabra") + ] + where + action = do + put "foo" "bar" + put "huga" "hoge" + put "abra" "cadabra" + multipleGet [ "foo" + , "huga" + , "unstored" + , "abra" + ] + +caseVsizRecord :: Assertion +caseVsizRecord = + action `returns` Right (Just 3, Nothing) + where + action = do + put "foo" "bar" + stored <- valueSize "foo" + unstored <- valueSize "bar" + return (stored, unstored) + +caseIterinit :: Assertion +caseIterinit = + action `returns` Right True + where + action = do + put "foo" "bar" + put "fuga" "hoge" + put "abra" "cadabra" + iterInit + key1 <- iterNext + _ <- iterNext + iterInit + key2 <- iterNext + return $ key1 == key2 + +caseIternext :: Assertion +caseIternext = + action `returns` Right [ Just "abra", Just "foo", Just "fuga" ] + where + action = do + put "foo" "bar" + put "fuga" "hoge" + put "abra" "cadabra" + iterInit + key1 <- iterNext + key2 <- iterNext + key3 <- iterNext + return $ sort [key1, key2, key3] + +caseIternextInvalid :: Assertion +caseIternextInvalid = + action `returns` Right Nothing + where + action = do + iterNext + +caseFwmkeys :: Assertion +caseFwmkeys = + action `returns` Right [ "abra", "abrac" ] + where + action = do + put "abr" "acadabra" + put "abra" "cadabra" + put "abrac" "adabra" + put "abraca" "dabra" + sort <$> forwardMatchingKeys "abra" 2