module DB (DB, Key(..), byJid, byNode, mk, get, getEnum, del, set, setEnum, sadd, srem, smembers, foldKeysM, hset, hdel, hgetall) where import Prelude () import BasicPrelude import Control.Error (readZ) import Network.Protocol.XMPP (JID(..), strNode) import qualified Database.TokyoCabinet as TC import qualified Data.Text as T import Util data DB = DB { tcdb :: TC.HDB } newtype Key = Key [String] openTokyoCabinet :: (TC.TCDB a) => String -> IO a openTokyoCabinet pth = TC.runTCM $ do db <- TC.new True <- TC.open db pth [TC.OREADER, TC.OWRITER, TC.OCREAT] return db mk :: String -> IO DB mk tcPath = do tcdb <- openTokyoCabinet tcPath return $ DB tcdb tcKey :: Key -> String tcKey (Key key) = intercalate "\0" key tcParseKey :: String -> Key tcParseKey str = Key $ map textToString $ T.split (=='\0') $ fromString str get :: DB -> Key -> IO (Maybe Text) get db key = fmap fromString <$> (TC.runTCM $ TC.get (tcdb db) $ tcKey key) getEnum :: (Enum a) => DB -> Key -> IO (Maybe a) getEnum db key = fmap toEnum <$> (TC.runTCM $ TC.get (tcdb db) $ tcKey key) del :: DB -> Key -> IO () del db key = do True <- TC.runTCM $ TC.out (tcdb db) $ tcKey key return () set :: DB -> Key -> Text -> IO () set db key val = do True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (textToString val) return () setEnum :: (Enum a) => DB -> Key -> a -> IO () setEnum db key val = do True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (fromEnum val) return () sadd :: DB -> Key -> [Text] -> IO () sadd db key new = do existing <- (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (show $ nub $ (map textToString new) ++ existing) return () srem :: DB -> Key -> [Text] -> IO () srem db key toremove = do existing <- (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (show $ filter (`notElem` toremove) existing) return () smembers :: (Read r) => DB -> Key -> IO [r] smembers db key = (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) hset :: (Eq k, Read k, Show k, Read v, Show v) => DB -> Key -> [(k, v)] -> IO () hset db key newitems = do items <- (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) let items' = nubBy (equating fst) (newitems ++ items) True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (show items') return () -- WARNING: Right now this function assumes all values are of type Maybe String hdel :: DB -> Key -> [Text] -> IO () hdel db key toremove = do items <- (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) let items' = filter ((`notElem` toremove) . fst) (items :: [(Text, Maybe String)]) True <- TC.runTCM $ TC.put (tcdb db) (tcKey key) (show items') return () hgetall :: (Read k, Read v) => DB -> Key -> IO [(k, v)] hgetall db key = (fromMaybe [] . (readZ =<<)) <$> TC.runTCM (TC.get (tcdb db) $ tcKey key) foldKeysM :: DB -> Key -> b -> (b -> Key -> IO b) -> IO b foldKeysM db (Key prefix) z f = do keys <- TC.runTCM $ TC.fwmkeys (tcdb db) (tcKey $ Key $ prefix ++ [""]) maxBound foldM f z $ map tcParseKey (keys :: [String]) byJid :: JID -> [String] -> Key byJid jid subkey = Key $ (textToString $ bareTxt jid) : subkey -- | Used when we know the JID is @cheogram.com, for example -- So usually this is ByTel, really byNode :: JID -> [String] -> Key byNode (JID { jidNode = Just node }) subkey = Key $ (textToString $ strNode node) : subkey byNode jid _ = error $ "JID without node used in byNode: " ++ show jid