etcd.py

import json
import logging
import pycurl
import re
import requests
import socket
import sys
import unittest

from StringIO import StringIO

class EtcdError(BaseException):
  #Generic etcd error
  pass

class Etcd(object):
  def __init__(self, logger, server=None):
    if server:
      self.server = server
    else:
      self.server = socket.gethostname()
    self.url = 'http://%(hostname)s:4001/v1/keys' % {
      'hostname': self.server}
    self.logger = logger

  def set_key(self, key, value):
    url = '%(base)s/%(key)s' % {
      'base': self.url,
      'key': key
    }
    data = 'value=%s' % value

    self.logger.debug("Saving data: %(data)s to %(url)s" %{
      'data': data,
      'url': url
    })
    storage = StringIO()

    curl = pycurl.Curl()
    curl.setopt(curl.URL, url)
    curl.setopt(curl.POSTFIELDS, data)
    curl.setopt(pycurl.FOLLOWLOCATION, 1)
    curl.setopt(pycurl.MAXREDIRS, 5)
    curl.setopt(curl.WRITEFUNCTION, storage.write)
    curl.perform()
    response = curl.getinfo(pycurl.HTTP_CODE)
    curl.close()

    if response == requests.codes.ok:
      return True
    else:
      self.logger.error("ETCD returned %(status)s %(text)s" % {
        'status': response,
        'text': storage.getvalue()})
      return None

  def get_key(self, key):
    url = '%(base)s/%(key)s' % {
      'base': self.url,
      'key': key
    }
    self.logger.debug('Getting url: ' + url)
    response = requests.get(url)
    self.logger.debug('Response: ' + response.text)

    res = json.loads(response.text)
    if isinstance(res, list):
      raise ValueError('Key "%s" is a directory, expecting leaf (use \
list_directory() to get directory listing).' % key)      

    #Check to see if Etcd returned an error
    if 'errorCode' in res:
      raise EtcdError(res['errorCode'], res['message']) 

    return str(res['value'])

  def delete_key(self, key):
    url = '%(base)s/%(key)s' % {
      'base': self.url,
      'key': key
    }

    response = requests.delete(url)
    if response.status_code == requests.codes.ok:
      return response.text
    else:
      response.raise_for_status()
      return None

  def list_directory(self, path):
    url = '%(base)s/%(path)s' % {
      'base': self.url,
      'path': path
    }
    response = requests.get(url)
    if response.status_code == requests.codes.ok:
      directory_list = []
      json_txt = json.loads(response.text)
      for entry in json_txt: 
        directory_list.append(str(entry['key']))
      return directory_list
    else:
      response.raise_for_status()
      return None

  def get_machines(self):
    url = '%(base)s/_etcd/machines' % {
      'base': self.url}
    res = json.loads(requests.get(url).text)

    #Check to see if Etcd returned an error
    if 'errorCode' in res:
      raise EtcdError(res['errorCode'], res['message']) 

    server_list = []
    for entry in res:
      server_list.append(str(entry['value']))

    return server_list

class TestEtcd(unittest.TestCase):
  def setUp(self):
    logger = logging.getLogger()
    stream = logging.StreamHandler(sys.stdout)
    stream.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    stream.setFormatter(formatter)
    logger.addHandler(stream)
    self.etcd = Etcd(logger)

  def test_a_setkey(self):
    ret = self.etcd.set_key('message', 'Hello World')
    self.assertTrue(ret)

  def test_b_getkey(self):
    self.etcd.set_key('message', 'Hello World')
    text = self.etcd.get_key('message')
    self.assertEqual(text, 'Hello World')

  def test_c_deletekey(self):
    #Set the key first before deleting it
    self.etcd.set_key('message', 'Hello World')

    text = self.etcd.delete_key('message')
    regex = re.compile(r'{"action":"DELETE","key":"/message",.*"index":\d+}')
    self.assertRegexpMatches(text, regex)

  def test_d_directorylist(self):
    #List a directory in Etcd
    dir_list = self.etcd.list_directory('formations/cholcomb')
    self.assertIsInstance(dir_list, list)