Merge pull request gmr#81 from fusiondog/restore-prune

Adding targets for ls, backup, and restore.  Adds pruning to restore.

Author: gmr

consulate/api/kv.py

@@ -225,7 +225,7 @@ def keys(self):
         """
         return sorted([row['Key'] for row in self._get_all_items()])
 
-    def records(self):
+    def records(self, key=None):
         """Return a list of tuples for all of the records in the Key/Value
         service
 
@@ -242,8 +242,12 @@ def records(self):
         :rtype: list of (Key, Flags, Value)
 
         """
-        return [(item['Key'], item['Flags'], item['Value'])
-                for item in self._get_all_items()]
+        if key:
+            return [(item['Key'], item['Flags'], item['Value'])
+                    for item in self._get_list([key], {'recurse': None})]
+        else:
+            return [(item['Key'], item['Flags'], item['Value'])
+                    for item in self._get_all_items()]
 
     def release_lock(self, item, session):
         """Release an existing lock from the Consul KV database.

consulate/cli.py

@@ -55,13 +55,19 @@ def connection_error():
 
 KV_PARSERS = [
     ('backup', 'Backup to stdout or a JSON file', [
+        [['key'], {'help': 'The key to use as target to backup a specific key or folder.', 
+                   'nargs':'?'}],
         [['-b', '--base64'], {'help': 'Base64 encode values',
                               'action': 'store_true'}],
         [['-f', '--file'], {'help': 'JSON file to write instead of stdout',
                             'nargs': '?'}],
         [['-p', '--pretty'], {'help': 'pretty-print JSON output',
                               'action': 'store_true'}]]),
     ('restore', 'Restore from stdin or a JSON file', [
+        [['key'], {'help': 'The key as target to restore to a specific key or folder.',
+                   'nargs':'?'}],
+        [['-p', '--prune'], {'help':'Remove entries from consul tree that are not in restore file.',
+                              'action': 'store_true'}],
         [['-b', '--base64'], {'help': 'Restore from Base64 encode values',
                               'action': 'store_true'}],
         [['-f', '--file'],
@@ -71,6 +77,8 @@ def connection_error():
          {'help': 'Do not replace existing entries',
           'action': 'store_true'}]]),
     ('ls', 'List all of the keys', [
+        [['key'], {'help': 'The key to use as target to list contents of specific key or folder',
+                   'nargs':'?'}],
         [['-l', '--long'],
          {'help': 'Long format',
           'action': 'store_true'}]]),
@@ -277,7 +285,13 @@ def kv_backup(consul, args):
 
     """
     handle = open(args.file, 'w') if args.file else sys.stdout
-    records = consul.kv.records()
+    if args.key:
+        args.key = args.key.strip('/')
+        prefixlen = len(args.key.split('/'))
+        records = [('/'.join(k.split('/')[prefixlen:]), f, v)
+                   for k, f, v in consul.kv.records(args.key)]
+    else:
+        records = consul.kv.records()
     if args.base64:
         if utils.PYTHON3:
             records = [(k, f, str(base64.b64encode(utils.maybe_encode(v)),
@@ -341,7 +355,12 @@ def kv_ls(consul, args):
 
     """
     try:
-        for key in consul.kv.keys():
+        if args.key:
+            args.key = args.key.lstrip('/')
+            keylist = sorted(consul.kv.find(args.key))
+        else:
+            keylist = consul.kv.keys()
+        for key in keylist:
             if args.long:
                 keylen = 0
                 if consul.kv[key]:
@@ -375,6 +394,12 @@ def kv_restore(consul, args):
     :param argparser.namespace args: The cli args
 
     """
+    if args.prune:
+        if args.key:
+            args.key = args.key.strip('/')
+            keylist = consul.kv.find(args.key)
+        else:
+            keylist = consul.kv.find('')
     handle = open(args.file, 'r') if args.file else sys.stdin
     data = json.load(handle)
     for row in data:
@@ -390,10 +415,28 @@ def kv_restore(consul, args):
         # Here's an awesome thing to make things work
         if not utils.PYTHON3 and isinstance(row[2], unicode):
             row[2] = row[2].encode('utf-8')
+        if args.key:
+            if row[0] == "":
+                rowkey = args.key
+            else:
+                rowkey = args.key + '/' + row[0]
+        else:
+            rowkey = row[0]
+        if args.prune:
+            if rowkey in keylist:
+                del keylist[rowkey]
         try:
-            consul.kv.set_record(row[0], row[1], row[2], not args.no_replace)
+            consul.kv.set_record(rowkey, row[1], row[2], not args.no_replace)
         except exceptions.ConnectionError:
             connection_error()
+    if args.prune:
+        for key in keylist:
+            print "Pruning " + key
+            try:
+                consul.kv.delete(key)
+            except exceptions.ConnectionError:
+                connection_error()
+            
 
 
 def kv_rm(consul, args):