blob: e69f441c896e83ce947897ed593f71b79161d6cf [file] [log] [blame]
Jordan Halterman00e92da2018-05-22 23:05:52 -07001#!/usr/bin/env python
2"""
3usage: atomix-gen-config [-h] [-s PARTITION_SIZE] [-n NUM_PARTITIONS]
4 [node_ip] [filename] [node_ip [node_ip ...]]
5
6Generate the partitions json file given a list of IPs or from the $OCC*
7environment variables.
8
9positional arguments:
10 filename File to write output to. If none is provided, output
11 is written to stdout.
12 node_ip IP Address(es) of the node(s) in the cluster. If no
13 IPs are given, will use the $OCC* environment
14 variables. NOTE: these arguemnts are only processed
15 after the filename argument.
16
17optional arguments:
18 -h, --help show this help message and exit
19 -s PARTITION_SIZE, --partition-size PARTITION_SIZE
20 Number of nodes per partition. Note that partition
21 sizes smaller than 3 are not fault tolerant. Defaults
22 to 3.
23 -n NUM_PARTITIONS, --num-partitions NUM_PARTITIONS
24 Number of partitions. Defaults to the number of nodes
25 in the cluster.
26"""
27
28from os import environ
29import argparse
30import re
31import json
32
33convert = lambda text: int(text) if text.isdigit() else text.lower()
34alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
35
Jordan Haltermand326de12018-07-31 11:57:21 -070036def get_vars_by_type(type):
37 vars = []
38 for var in environ:
39 if re.match(r"{}[0-9]+".format(type), var):
40 vars.append(var)
41 return sorted(vars, key=alphanum_key)
42
43
44def get_vars():
45 vars = get_vars_by_type('OCC')
46 if len(vars) == 0:
47 vars = get_vars_by_type('OC')
48 return vars
49
Jordan Halterman00e92da2018-05-22 23:05:52 -070050
51def get_local_node(node, ips=None):
52 if not ips:
Jordan Haltermand326de12018-07-31 11:57:21 -070053 ips = [environ[v] for v in get_vars()]
Jordan Halterman00e92da2018-05-22 23:05:52 -070054 return 'atomix-{}'.format(ips.index(node) + 1)
55
Jordan Haltermand326de12018-07-31 11:57:21 -070056
Jordan Halterman00e92da2018-05-22 23:05:52 -070057def get_nodes(ips=None, default_port=5679):
58 node = lambda id, ip, port: {'id': id, 'address': '{}:{}'.format(ip, port)}
59 result = []
60 if not ips:
Jordan Haltermand326de12018-07-31 11:57:21 -070061 ips = [environ[v] for v in get_vars()]
Jordan Halterman00e92da2018-05-22 23:05:52 -070062 i = 1
63 for ip_string in ips:
64 address_tuple = ip_string.split(":")
65 if len(address_tuple) == 3:
66 id=address_tuple[0]
67 ip=address_tuple[1]
68 port=int(address_tuple[2])
69 else:
70 id='atomix-{}'.format(i)
71 i += 1
72 ip=ip_string
73 port=default_port
74 result.append(node(id, ip, port))
75 return result
76
Jordan Haltermand326de12018-07-31 11:57:21 -070077
Jordan Halterman00e92da2018-05-22 23:05:52 -070078def get_local_address(node, ips=None, default_port=5679):
Jordan Halterman00e92da2018-05-22 23:05:52 -070079 if not ips:
Jordan Haltermand326de12018-07-31 11:57:21 -070080 ips = [environ[v] for v in get_vars()]
Jordan Halterman00e92da2018-05-22 23:05:52 -070081 for ip_string in ips:
82 address_tuple = ip_string.split(":")
83 if len(address_tuple) == 3:
84 id=address_tuple[0]
85 ip=address_tuple[1]
86 port=int(address_tuple[2])
87 if node == id or node == ip:
88 return '{}:{}'.format(ip, port)
89 return '{}:{}'.format(node, default_port)
90
91if __name__ == '__main__':
92 parser = argparse.ArgumentParser(
93 description="Generate the partitions json file given a list of IPs or from the $OCC* environment variables.")
94 parser.add_argument(
95 '-s', '--partition-size', type=int, default=3,
96 help="Number of nodes per partition. Note that partition sizes smaller than 3 are not fault tolerant. Defaults to 3." )
97 parser.add_argument(
98 '-n', '--num-partitions', type=int,
99 help="Number of partitions. Defaults to the number of nodes in the cluster." )
100 # TODO: make filename and nodes independent. This will break backwards compatibility with existing usage.
101 parser.add_argument(
102 'node', metavar='node_ip', type=str, help='IP address of the node for which to generate the configuration')
103 parser.add_argument(
104 'filename', metavar='filename', type=str, nargs='?',
105 help='File to write output to. If none is provided, output is written to stdout.')
106 parser.add_argument(
107 'nodes', metavar='node_ip', type=str, nargs='*',
108 help='IP Address(es) of the node(s) in the cluster. If no IPs are given, ' +
109 'will use the $OCC* environment variables. NOTE: these arguemnts' +
110 ' are only processed after the filename argument.')
111
112 args = parser.parse_args()
113 filename = args.filename
114 partition_size = args.partition_size
115 local_member_id = get_local_node(args.node)
116 local_member_address = get_local_address(args.node, args.nodes)
117 nodes = get_nodes(args.nodes)
118 num_partitions = args.num_partitions
119 if not num_partitions:
120 num_partitions = len(nodes)
121
122 data = {
123 'cluster': {
124 'clusterId': 'onos',
125 'node': {
126 'id': local_member_id,
127 'address': local_member_address
128 },
129 'discovery': {
130 'type': 'bootstrap',
131 'nodes': nodes
132 }
133 },
134 'managementGroup': {
135 'type': 'raft',
136 'partitions': 1,
137 'partitionSize': len(nodes),
138 'members': [node['id'] for node in nodes]
139 },
140 'partitionGroups': {
141 'raft': {
142 'type': 'raft',
143 'partitions': num_partitions,
144 'partitionSize': partition_size,
145 'members': [node['id'] for node in nodes]
146 }
147 }
148 }
149 output = json.dumps(data, indent=4)
150
151 if filename:
152 with open(filename, 'w') as f:
153 f.write(output)
154 else:
155 print output