Merge pull request #30627 from sjmh/etcd_set_tree

Add etcd update function

salt/modules/etcd_mod.py

@@ -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

salt/utils/etcd_util.py

@@ -187,6 +187,37 @@ class EtcdClient(object):
             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)
 

tests/unit/modules/etcd_mod_test.py

@@ -85,6 +85,36 @@ class EtcdModTestCase(TestCase):
             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):

tests/unit/utils/etcd_util_test.py

@@ -188,6 +188,91 @@ class EtcdUtilTestCase(TestCase):
         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