Skip to content

Commit

Permalink
Add etcd update function
Browse files Browse the repository at this point in the history
Adds a update function that takes a dict and sets multiple etcd keys.
  • Loading branch information
Steve Hajducko committed Jan 25, 2016
1 parent a48ffaa commit a015473
Show file tree
Hide file tree
Showing 4 changed files with 199 additions and 0 deletions.
53 changes: 53 additions & 0 deletions salt/modules/etcd_mod.py
Expand Up @@ -101,6 +101,59 @@ def set_(key, value, profile=None, ttl=None, directory=False):
return client.set(key, value, ttl=ttl, directory=directory)


def update(fields, path='', profile=None):
'''
.. versionadded:: Boron
Sets a dictionary of values in one call. Useful for large updates
in syndic environments. The dictionary can contain a mix of formats
such as:
.. code-block:: python
{
'/some/example/key': 'bar',
'/another/example/key': 'baz'
}
Or it may be a straight dictionary, which will be flattened to look
like the above format:
.. code-block:: python
{
'some': {
'example': {
'key': 'bar'
}
},
'another': {
'example': {
'key': 'baz'
}
}
}
You can even mix the two formats and it will be flattened to the first
format. Leading and trailing '/' will be removed.
Empty directories can be created by setting the value of the key to an
empty dictionary.
The 'path' parameter will optionally set the root of the path to use.
CLI Example:
.. code-block:: bash
salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}"
salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}" profile=my_etcd_config
salt myminion etcd.update "{'/path/to/key': 'baz', '/another/key': 'bar'}" path='/some/root'
'''
client = __utils__['etcd_util.get_conn'](__opts__, profile)
return client.update(fields, path)


def watch(key, recurse=False, profile=None, timeout=0, index=None):
'''
.. versionadded:: Boron
Expand Down
31 changes: 31 additions & 0 deletions salt/utils/etcd_util.py
Expand Up @@ -187,6 +187,37 @@ def read(self, key, recursive=False, wait=False, timeout=None, waitIndex=None):
raise
return result

def _flatten(self, data, path=''):
if len(data.keys()) == 0:
return {path: {}}
path = path.strip('/')
flat = {}
for k, v in data.iteritems():
k = k.strip('/')
if path:
p = '/{0}/{1}'.format(path, k)
else:
p = '/{0}'.format(k)
if isinstance(v, dict):
ret = self._flatten(v, p)
flat.update(ret)
else:
flat[p] = v
return flat

def update(self, fields, path=''):
if not isinstance(fields, dict):
log.error('etcd.update: fields is not type dict')
return None
fields = self._flatten(fields, path)
keys = {}
for k, v in fields.iteritems():
is_dir = False
if isinstance(v, dict):
is_dir = True
keys[k] = self.write(k, v, directory=is_dir)
return keys

def set(self, key, value, ttl=None, directory=False):
return self.write(key, value, ttl=ttl, directory=directory)

Expand Down
30 changes: 30 additions & 0 deletions tests/unit/modules/etcd_mod_test.py
Expand Up @@ -85,6 +85,36 @@ def test_set(self):
self.instance.set.side_effect = Exception
self.assertRaises(Exception, etcd_mod.set_, 'err', 'stack')

# 'update' function tests: 1

def test_update(self):
'''
Test if can set multiple keys in etcd
'''
with patch.dict(etcd_mod.__utils__, {'etcd_util.get_conn': self.EtcdClientMock}):
args = {
'x': {
'y': {
'a': '1',
'b': '2',
}
},
'z': '4',
'd': {},
}

result = {
'/some/path/x/y/a': '1',
'/some/path/x/y/b': '2',
'/some/path/z': '4',
'/some/path/d': {},
}
self.instance.update.return_value = result
self.assertDictEqual(etcd_mod.update(args, path='/some/path'), result)
self.instance.update.assert_called_with(args, '/some/path')
self.assertDictEqual(etcd_mod.update(args), result)
self.instance.update.assert_called_with(args, '')

# 'ls_' function tests: 1

def test_ls(self):
Expand Down
85 changes: 85 additions & 0 deletions tests/unit/utils/etcd_util_test.py
Expand Up @@ -188,6 +188,91 @@ def test_write(self, mock):
etcd_client.write.side_effect = Exception
self.assertRaises(Exception, client.set, 'some-key', 'some-val')

@patch('etcd.Client', autospec=True)
def test_flatten(self, mock):
client = etcd_util.EtcdClient({})
some_data = {
'/x/y/a': '1',
'x': {
'y': {
'b': '2'
}
},
'm/j/': '3',
'z': '4',
'd': {},
}

result_path = {
'/test/x/y/a': '1',
'/test/x/y/b': '2',
'/test/m/j': '3',
'/test/z': '4',
'/test/d': {},
}

result_nopath = {
'/x/y/a': '1',
'/x/y/b': '2',
'/m/j': '3',
'/z': '4',
'/d': {},
}

result_root = {
'/x/y/a': '1',
'/x/y/b': '2',
'/m/j': '3',
'/z': '4',
'/d': {},
}

self.assertEqual(client._flatten(some_data, path='/test'), result_path)
self.assertEqual(client._flatten(some_data, path='/'), result_root)
self.assertEqual(client._flatten(some_data), result_nopath)

@patch('etcd.Client', autospec=True)
def test_update(self, mock):
client = etcd_util.EtcdClient({})
some_data = {
'/x/y/a': '1',
'x': {
'y': {
'b': '3'
}
},
'm/j/': '3',
'z': '4',
'd': {},
}

result = {
'/test/x/y/a': '1',
'/test/x/y/b': '2',
'/test/m/j': '3',
'/test/z': '4',
'/test/d': True,
}

flatten_result = {
'/test/x/y/a': '1',
'/test/x/y/b': '2',
'/test/m/j': '3',
'/test/z': '4',
'/test/d': {}
}
client._flatten = MagicMock(return_value=flatten_result)

self.assertEqual(client.update('/some/key', path='/blah'), None)

with patch.object(client, 'write', autospec=True) as write_mock:
def write_return(key, val, ttl=None, directory=None):
return result.get(key, None)
write_mock.side_effect = write_return
self.assertDictEqual(client.update(some_data, path='/test'), result)
client._flatten.assert_called_with(some_data, '/test')
self.assertEqual(write_mock.call_count, 5)

@patch('etcd.Client', autospec=True)
def test_rm(self, mock):
etcd_client = mock.return_value
Expand Down

0 comments on commit a015473

Please sign in to comment.