Merge "Adds group support to inventory-manage.py"

scripts/inventory-manage.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+#
 # Copyright 2014, Rackspace US, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,6 +15,9 @@
 # limitations under the License.
 #
 # (c) 2014, Kevin Carter <kevin.carter@rackspace.com>
+# (c) 2015, Major Hayden <major@mhtx.net>
+#
+"""Returns data about containers and groups in tabular formats."""
 import argparse
 import json
 import os
@@ -43,9 +47,9 @@ def file_find(filename, user_file=None, pass_exception=False):
     if user_file is not None:
         file_check.insert(0, os.path.expanduser(user_file))
 
-    for f in file_check:
+    for filename in file_check:
-        if os.path.isfile(f):
+        if os.path.isfile(filename):
-            return f
+            return filename
     else:
         if pass_exception is False:
             raise SystemExit('No file found at: %s' % file_check)
@@ -54,6 +58,12 @@ def file_find(filename, user_file=None, pass_exception=False):
 
 
 def recursive_list_removal(inventory, purge_list):
+    """Remove items from a list.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    purge_list -- list of items to remove
+    """
     for item in purge_list:
         for _item in inventory:
             if item == _item:
@@ -61,6 +71,12 @@ def recursive_list_removal(inventory, purge_list):
 
 
 def recursive_dict_removal(inventory, purge_list):
+    """Remove items from a dictionary.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    purge_list -- list of items to remove
+    """
     for key, value in inventory.iteritems():
         if isinstance(value, dict):
             for _key, _value in value.iteritems():
@@ -112,11 +128,143 @@ def args():
         action='store_true',
         default=False
     )
+    exclusive_action.add_argument(
+        '-g',
+        '--list-groups',
+        help='List groups and containers in each group',
+        action='store_true',
+        default=False
+    )
+    exclusive_action.add_argument(
+        '-G',
+        '--list-containers',
+        help='List containers and their groups',
+        action='store_true',
+        default=False
+    )
 
     return vars(parser.parse_args())
 
 
+def get_all_groups(inventory):
+    """Retrieve all ansible groups.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+
+    Will return a dictionary of containers as keys and corresponding groups
+    as values.
+    """
+    containers = {}
+    for container_name in inventory['_meta']['hostvars'].keys():
+
+        # Skip the default group names since they're not helpful (like aio1).
+        if '_' not in container_name:
+            continue
+
+        groups = get_groups_for_container(inventory, container_name)
+        containers[container_name] = groups
+
+    return containers
+
+
+def get_groups_for_container(inventory, container_name):
+    """Return groups for a particular container.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    container_name -- name of a container to lookup
+
+    Will return a list of groups that the container belongs to.
+    """
+    # Beware, this dictionary comprehension requires Python 2.7, but we should
+    # have this on openstack-ansible hosts already.
+    groups = {k for (k, v) in inventory.items() if
+              ('hosts' in v
+              and container_name in v['hosts'])}
+    return groups
+
+
+def get_containers_for_group(inventory, group):
+    """Return containers that belong to a particular group.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    group -- group to use to lookup containers
+
+    Will return a list of containers that belong to a group, or None if no
+    containers match the group provided.
+    """
+    if 'hosts' in inventory[group]:
+        containers = inventory[group]['hosts']
+    else:
+        containers = None
+    return containers
+
+
+def print_groups_per_container(inventory):
+    """Return a table of containers and the groups they belong to.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    """
+    containers = get_all_groups(inventory)
+    required_list = [
+        'container_name',
+        'groups'
+    ]
+    table = prettytable.PrettyTable(required_list)
+
+    for container_name, groups in containers.iteritems():
+        row = [container_name, ', '.join(sorted(groups))]
+        table.add_row(row)
+
+    for tbl in table.align.keys():
+        table.align[tbl] = 'l'
+
+    return table
+
+
+def print_containers_per_group(inventory):
+    """Return a table of groups and the containers in each group.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    """
+    required_list = [
+        'groups',
+        'container_name'
+    ]
+    table = prettytable.PrettyTable(required_list)
+
+    for group_name in inventory.keys():
+        containers = get_containers_for_group(inventory, group_name)
+
+        # Don't show a group if it has no containers
+        if containers is None or len(containers) < 1:
+            continue
+
+        # Don't show default group
+        if len(containers) == 1 and '_' not in containers[0]:
+            continue
+
+        # Join with newlines here to avoid having a horrific table with tons
+        # of line wrapping.
+        row = [group_name, '\n'.join(containers)]
+        table.add_row(row)
+
+    for tbl in table.align.keys():
+        table.align[tbl] = 'l'
+
+    return table
+
+
 def print_inventory(inventory, sort_key):
+    """Return a table of containers with detail about each.
+
+    Keyword arguments:
+    inventory -- inventory dictionary
+    """
     _meta_data = inventory['_meta']['hostvars']
     required_list = [
         'container_name',
@@ -155,15 +303,25 @@ def main():
 
     # Get the contents of the system environment json
     environment_file = file_find(filename=user_args['file'])
-    with open(environment_file, 'rb') as f:
+    with open(environment_file, 'rb') as f_handle:
-        inventory = json.loads(f.read())
+        inventory = json.loads(f_handle.read())
 
+    # Make a table with hosts in the left column and details about each in the
+    # columns to the right
     if user_args['list_host'] is True:
         print(print_inventory(inventory, user_args['sort']))
+
+    # Groups in first column, containers in each group on the right
+    elif user_args['list_groups'] is True:
+        print(print_groups_per_container(inventory))
+
+    # Containers in the first column, groups for each container on the right
+    elif user_args['list_containers'] is True:
+        print(print_containers_per_group(inventory))
     else:
         recursive_dict_removal(inventory, user_args['remove_item'])
-        with open(environment_file, 'wb') as f:
+        with open(environment_file, 'wb') as f_handle:
-            f.write(json.dumps(inventory, indent=2))
+            f_handle.write(json.dumps(inventory, indent=2))
         print('Success. . .')
 
 if __name__ == "__main__":