~singpolyma/cheogram

ref: 6349966097b44cc4a1f32f12f2eefae75d6b95c3 cheogram/DB.hs -rw-r--r-- 3.5 KiB
63499660Stephen Paul Weber Abstract DB 1 year, 3 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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