Refactor etcd updating into its own class.

report.py

@@ -1,55 +1,8 @@
 #!/usr/bin/env python3
 
-from collections import defaultdict
 from monitor import Monitor
+from updater import Updater
 import argparse
-import etcd
-
-
-def etcd_put(client, prefix, obj):
-  for key, value in obj.items():
-    new_prefix = "%s/%s" % (prefix, key)
-
-    if isinstance(value, dict):
-      etcd_put(client, new_prefix, value)
-    else:
-      client.write(new_prefix, str(value))
-
-
-def add_containers(new_containers):
-  global containers, host_index, label_index, network_index
-  for container in new_containers:
-    name = container['name']
-    containers[name] = container
-
-    for k, v in container['labels'].items():
-      label_index[k][name] = v
-    for k, v in container['net']['addr'].items():
-      network_index[k][name] = v
-
-    host_index[host][name] = name
-
-  etcd_put(etcd_client, prefix + '/containers', containers)
-  etcd_put(etcd_client, prefix + '/labels', label_index)
-  etcd_put(etcd_client, prefix + '/networks', network_index)
-  etcd_put(etcd_client, prefix + '/hosts', host_index)
-
-
-def remove_containers(old_containers):
-  global containers, host_index, label_index, network_index
-  for container in old_containers:
-    name = container['name']
-    del containers[name]
-    etcd_client.delete(prefix + '/containers/' + name, recursive=True)
-
-    for k, v in container['labels'].items():
-      del label_index[k][name]
-      etcd_client.delete(prefix + '/labels/' + k + '/' + name)
-    for k, v in container['net']['addr'].items():
-      del network_index[k][name]
-      etcd_client.delete(prefix + '/networks/' + k + '/' + name)
-    etcd_client.delete(prefix + '/hosts/' + host + '/' + name)
-
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--name', help='Name of this docker host', default='unknown')
@@ -58,19 +11,9 @@ parser.add_argument('--etcd-host', help='Host to connect to etcd on', default='e
 parser.add_argument('--etcd-prefix', help='Prefix to use when adding keys to etcd', default='/docker')
 args = parser.parse_args()
 
-monitor = Monitor(args.name, add_containers, remove_containers)
-etcd_client = etcd.Client(host=args.etcd_host, port=args.etcd_port)
-prefix = args.etcd_prefix
-host = args.name
-
-containers = {}
-label_index = defaultdict(dict) 
-network_index = defaultdict(dict)
-host_index = defaultdict(dict)
-
-try:
-  etcd_client.delete(prefix, recursive=True)
-except etcd.EtcdKeyNotFound:
-  pass
+updater = Updater(args.etcd_host, args.etcd_port, args.etcd_prefix)
+monitor = Monitor(args.name, updater.add_containers, updater.remove_containers)
 
+updater.wipe()
 monitor.monitor()
+

updater.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+
+import etcd
+
+
+class Updater:
+
+  def __init__(self, host, port, prefix):
+    self._client = etcd.Client(host=host, port=port)
+    self._prefix = prefix
+
+
+  def _write(self, key, value):
+    self._client.write(self._prefix + key, value)
+
+
+  def _write_obj(self, prefix, obj):
+    for key, value in obj.items():
+      new_prefix = "%s/%s" % (prefix, key)
+
+      if isinstance(value, dict):
+        self._write_obj(new_prefix, value)
+      else:
+        self._write(new_prefix, str(value))
+
+
+  def _delete(self, key):
+    try:
+      self._client.delete(self._prefix + key, recursive=True)
+    except etcd.EtcdKeyNotFound:
+      pass
+
+
+  def wipe(self):
+    self._delete('')
+
+
+  def add_containers(self, new_containers):
+    for container in new_containers:
+      name = container['name']
+      self._write_obj('/containers/%s' % name, container)
+      for k, v in container['labels'].items():
+        self._write('/labels/%s/%s' % (k, name), v)
+      for k, v in container['net']['addr'].items():
+        self._write('/networks/%s/%s' % (k, name), v)
+      self._write('/hosts/%s/%s' % (container['host'], name), name)
+
+
+  def remove_containers(self, old_containers):
+    for container in old_containers:
+      name = container['name']
+      self._delete('/containers/%s' % name)
+      for k, v in container['labels'].items():
+        self._delete('/labels/%s/%s' % (k, name))
+      for k, v in container['net']['addr'].items():
+        self._delete('/networks/%s/%s' % (k, name)) 
+      self._delete('/hosts/%s/%s' % (container['host'], name))
+