forked from couchbaselabs/mcsoda
/
mcsoda_mongo.py
executable file
·82 lines (67 loc) · 2.57 KB
/
mcsoda_mongo.py
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
#!/usr/bin/env python
import sys
import copy
import math
import time
import mcsoda
import pymongo # Use: pip install pymongo==2.0
from pymongo import Connection
mongoDocCache = {}
class StoreMongo(mcsoda.Store):
def connect(self, target, user, pswd, cfg, cur):
self.cfg = cfg
self.cur = cur
self.target = target
self.host_port = (target + ":27017").split(':')[0:2]
self.host_port[1] = int(self.host_port[1])
self.conn = Connection(self.host_port[0],
self.host_port[1])
self.db = self.conn['default']
self.coll = self.db['default']
self.xfer_sent = 0
self.xfer_recv = 0
def gen_doc(self, key_num, key_str, min_value_size, json=None, cache=None):
global mongoDocCache
if json is None:
json = self.cfg.get('json', 1) > 0
if cache is None:
cache = self.cfg.get('doc-cache', 0)
d = None
if cache:
d = mongoDocCache.get(key_num, None)
if d is None:
d = { "_id": key_str,
"key_num": key_num,
"name": mcsoda.key_to_name(key_num, key_str),
"email": mcsoda.key_to_email(key_num, key_str),
"city": mcsoda.key_to_city(key_num, key_str),
"country": mcsoda.key_to_country(key_num, key_str),
"realm": mcsoda.key_to_realm(key_num, key_str),
"coins": mcsoda.key_to_coins(key_num, key_str),
"achievements": mcsoda.key_to_achievements(key_num, key_str) }
if cache:
mongoDocCache[key_num] = d
d = copy.deepcopy(d)
d['body'] = self.cfg['body'][min_value_size]
return d
def command(self, c):
cmd, key_num, key_str, data, expiration = c
cmd_start = time.time()
if cmd[0] == 'g':
self.coll.find_one(key_str)
elif cmd[0] == 's':
self.coll.save(data)
elif cmd[0] == 'd':
self.coll.remove(key_str)
else:
raise Exception("StoreMongo saw an unsupported cmd: " + cmd)
cmd_end = time.time()
histo = self.add_timing_sample(cmd, cmd_end - cmd_start)
if self.sc:
p = self.histo_percentile(histo, [0.90, 0.95, 0.99])
self.sc.latency_stats(cmd, p)
if __name__ == "__main__":
if sys.argv[1].find("mongo") != 0:
raise Exception("usage: mcsoda_mongo mongo://HOST:27017 ...")
argv = (' '.join(sys.argv) + ' doc-gen=0').split(' ')
mcsoda.main(argv, protocol="mongo", stores=[StoreMongo()])