1cffa2d40SAaron Miller#!/usr/bin/env python
2cffa2d40SAaron Millerimport os, sys, getopt, urllib2, urllib, json
3cffa2d40SAaron Miller
4cffa2d40SAaron Millerfrom urllib2 import HTTPError
5cffa2d40SAaron Miller
6cffa2d40SAaron Millervalid_bucket_types = ["membase", "memcached"]
7b420b4a8SAliaksey Artamonauvalid_service_types = {"kv", "n1ql", "index", "fts", "example"}
8cffa2d40SAaron Miller
9cffa2d40SAaron Millerdef usage():
10cffa2d40SAaron Miller    print "usage: \n\
11cffa2d40SAaron Miller            -n <number of nodes>\n\
126e56d478SArtem Stemkovski            -T <services to run. kv if unspecified> (eg: n0:kv,n1:index+n1ql+fts)\n\
13cffa2d40SAaron Miller            -s <memory size (min 256) default: 256\n\
14f881138cSAliaksey Artamonau            -I <index memory size> default: 256\n\
152bf87e3cSDale Harvey            -t <bucket type> (membase, memcached) default: membase\n\
16715eb21cSAliaksey Artamonau            -r <num replicas> (max 3) default: 1 (Only for membase buckets!)\n\
1734d72ac4SAliaksey Artamonau            -i (don't index replicas) default: replica index enabled\n\
1834d72ac4SAliaksey Artamonau            -S <start index> default: 0"
19cffa2d40SAaron Miller
20cffa2d40SAaron Miller
21cffa2d40SAaron Millerclass PasswordManager(urllib2.HTTPPasswordMgr):
22cffa2d40SAaron Miller    def __init__(self, username, password):
23cffa2d40SAaron Miller        self.auth = (username, password)
24cffa2d40SAaron Miller
25cffa2d40SAaron Miller    def find_user_password(self, realm, authuri):
26cffa2d40SAaron Miller        return self.auth
27cffa2d40SAaron Miller
28cffa2d40SAaron Miller
29cffa2d40SAaron Millerdef main():
30cffa2d40SAaron Miller    try:
3186ae514bSAliaksey Artamonau        opts, args = getopt.getopt(sys.argv[1:],
3286ae514bSAliaksey Artamonau                                   "n:t:s:r:iS:T:I:", ["dont-rebalance"])
33cffa2d40SAaron Miller    except getopt.GetoptError, err:
34cffa2d40SAaron Miller        print str(err)
35cffa2d40SAaron Miller        usage()
36cffa2d40SAaron Miller        sys.exit()
37cffa2d40SAaron Miller    nodes = 0
382bf87e3cSDale Harvey    buckettype = "membase"
39cffa2d40SAaron Miller    memsize = 256
40f881138cSAliaksey Artamonau    indexmemsize = 256
4148c06e1bSAliaksey Artamonau    replicas = 1
42715eb21cSAliaksey Artamonau    replica_index = True
4334d72ac4SAliaksey Artamonau    start_index = 0
44cf1212edSAliaksey Artamonau    deploy = ['kv']
4586ae514bSAliaksey Artamonau    do_rebalance = True
46715eb21cSAliaksey Artamonau
47cffa2d40SAaron Miller    data_base_path = os.getcwd() + "/data"
48cffa2d40SAaron Miller    for o, a in opts:
49cffa2d40SAaron Miller        if o == "-n":
50f547169dSSriram Melkote            nodes = int(a)
51cffa2d40SAaron Miller        elif o == "-t":
52cffa2d40SAaron Miller            buckettype = a
53cffa2d40SAaron Miller        elif o == "-s":
54cffa2d40SAaron Miller            memsize = a
55f881138cSAliaksey Artamonau        elif o == "-I":
56f881138cSAliaksey Artamonau            indexmemsize = a
57cffa2d40SAaron Miller        elif o == "-r":
58cffa2d40SAaron Miller            replicas = a
59715eb21cSAliaksey Artamonau        elif o == "-i":
60715eb21cSAliaksey Artamonau            replica_index = False
6134d72ac4SAliaksey Artamonau        elif o == "-S":
6234d72ac4SAliaksey Artamonau            start_index = int(a)
63184618c3SSriram Melkote        elif o == "-T":
64cf1212edSAliaksey Artamonau            plan = a.replace(' ','').split(',')
65cf1212edSAliaksey Artamonau
66cf1212edSAliaksey Artamonau            if len(plan) == 1 and len(plan[0].split(':')) == 1:
67cf1212edSAliaksey Artamonau                deploy = plan[0].split('+')
68cf1212edSAliaksey Artamonau            else:
69cf1212edSAliaksey Artamonau                plan = dict(e.split(':') for e in plan)
70cf1212edSAliaksey Artamonau                deploy = dict([(k, v.split('+')) for k, v in plan.items()])
7186ae514bSAliaksey Artamonau        elif o == "--dont-rebalance":
7286ae514bSAliaksey Artamonau            do_rebalance = False
73cffa2d40SAaron Miller        else:
74cffa2d40SAaron Miller            usage()
75cffa2d40SAaron Miller            sys.exit()
76f547169dSSriram Melkote
77cf1212edSAliaksey Artamonau    if isinstance(deploy, list):
78cf1212edSAliaksey Artamonau        services = deploy
79cf1212edSAliaksey Artamonau        deploy = dict(("n%d" % i, services[:]) for i in xrange(nodes))
80cf1212edSAliaksey Artamonau
81f547169dSSriram Melkote    deploy["n0"] = deploy.get("n0", []) + ["kv"]
82cffa2d40SAaron Miller    if nodes == 0 or buckettype not in valid_bucket_types or \
83184618c3SSriram Melkote            int(memsize) < 256 or int(replicas) > 3 or \
84f547169dSSriram Melkote            not set(deploy.keys()) <= set(["n" + str(i) for i in range(nodes)]) or \
85f547169dSSriram Melkote            not set(reduce(lambda x,y:x+y, deploy.values(), [])) <= valid_service_types:
86cffa2d40SAaron Miller        usage()
87cffa2d40SAaron Miller        sys.exit()
88cffa2d40SAaron Miller
89cffa2d40SAaron Miller    password_mgr = PasswordManager("Administrator", "asdasd")
90cffa2d40SAaron Miller    handler = urllib2.HTTPBasicAuthHandler(password_mgr)
91cffa2d40SAaron Miller    o = urllib2.build_opener(handler)
92cffa2d40SAaron Miller
9334d72ac4SAliaksey Artamonau    print "Connecting {0} nodes, bucket type {1}, mem size {2} " \
94f547169dSSriram Melkote        "with {3} replica copies, start index {4}, password asdasd, "\
95f547169dSSriram Melkote        "deployment plan {5}\n".format(nodes, buckettype,
96184618c3SSriram Melkote                                   memsize, replicas, start_index,
97f547169dSSriram Melkote                                   str(deploy))
9834d72ac4SAliaksey Artamonau
9934d72ac4SAliaksey Artamonau    base_port = 9000 + start_index
10034d72ac4SAliaksey Artamonau
101f547169dSSriram Melkote    services = deploy["n0"]
102f547169dSSriram Melkote    print "Connecting node 0 with services {0}".format(str(services))
103f547169dSSriram Melkote    o.open("http://127.0.0.1:{0}/node/controller/setupServices".format(base_port),
104f547169dSSriram Melkote           "services={0}".format(",".join(services))).read()
10534d72ac4SAliaksey Artamonau    o.open("http://127.0.0.1:{0}/pools/default".format(base_port),
106f881138cSAliaksey Artamonau           "memoryQuota=" + str(memsize) +
107f881138cSAliaksey Artamonau           "&indexMemoryQuota=" + str(indexmemsize)).read()
108a3cae4a3SArtem Stemkovski    o.open("http://127.0.0.1:{0}/pools/default/buckets".format(base_port),
109a3cae4a3SArtem Stemkovski           "name=default" +
110a3cae4a3SArtem Stemkovski           "&authType=sasl" +
111a3cae4a3SArtem Stemkovski           "&saslPassword=" +
112a3cae4a3SArtem Stemkovski           "&bucketType=" + buckettype +
11334d72ac4SAliaksey Artamonau           "&ramQuotaMB=" + str(memsize) +
11434d72ac4SAliaksey Artamonau           "&replicaNumber=" + str(replicas) +
11534d72ac4SAliaksey Artamonau           "&replicaIndex=" + bool_request_value(replica_index)).read()
11634d72ac4SAliaksey Artamonau    o.open("http://127.0.0.1:{0}/settings/web".format(base_port),
11734d72ac4SAliaksey Artamonau           "port=SAME&username=Administrator&password=asdasd").read()
118cffa2d40SAaron Miller
119f547169dSSriram Melkote    for i in range(1, nodes):
12034d72ac4SAliaksey Artamonau        port = base_port + i
121f547169dSSriram Melkote        services = deploy.get("n" + str(i), [])
122f547169dSSriram Melkote        if not services: services = ["kv"]
123f547169dSSriram Melkote        print "Connecting node {0} with services {1}".format(i, str(services))
124cffa2d40SAaron Miller        o.open("http://127.0.0.1:{0}/node/controller/doJoinCluster".format(port),
125184618c3SSriram Melkote               "user=Administrator&password=asdasd&" +
126184618c3SSriram Melkote                   "clusterMemberHostIp=127.0.0.1" +
127184618c3SSriram Melkote                   "&clusterMemberPort={0}".format(base_port) +
128184618c3SSriram Melkote                   "&services={0}".format(",".join(services))).read()
12934d72ac4SAliaksey Artamonau
13086ae514bSAliaksey Artamonau    if do_rebalance:
13186ae514bSAliaksey Artamonau        print "Getting node list"
13286ae514bSAliaksey Artamonau        info = json.loads(o.open("http://127.0.0.1:{0}/nodeStatuses".format(base_port)).read())
13334d72ac4SAliaksey Artamonau
13486ae514bSAliaksey Artamonau        print "Servers added, triggering rebalance."
13586ae514bSAliaksey Artamonau        o.open("http://127.0.0.1:{0}/controller/rebalance".format(base_port),
13686ae514bSAliaksey Artamonau               urllib.urlencode(
13786ae514bSAliaksey Artamonau                {'knownNodes': ",".join([info[k]['otpNode'] for k in info]),
13886ae514bSAliaksey Artamonau                 'ejectedNodes': ''})).read()
139cffa2d40SAaron Miller
140715eb21cSAliaksey Artamonaudef bool_request_value(Value):
141715eb21cSAliaksey Artamonau    return "1" if Value else "0"
142715eb21cSAliaksey Artamonau
143cffa2d40SAaron Millerif __name__ == '__main__':
144cffa2d40SAaron Miller    main()
145