Updates to ohi to pull cache if specified. Also require version

2 files changed, 66 insertions, 40 deletions

diff --git a/bin/ohi b/bin/ohi
index d679edcfb..be9c53ec0 100755
--- a/bin/ohi
+++ b/bin/ohi
@@ -1,14 +1,16 @@
 #!/usr/bin/env python
+'''
+Ohi = Openshift Host Inventory
+
+This script provides an easy way to look at your host inventory.
+
+This depends on multi_inventory being setup correctly.
+'''
 # vim: expandtab:tabstop=4:shiftwidth=4
 
 import argparse
-import traceback
 import sys
 import os
-import re
-import tempfile
-import time
-import subprocess
 import ConfigParser
 
 from openshift_ansible import awsutil
@@ -20,6 +22,9 @@ CONFIG_HOST_TYPE_ALIAS_SECTION = 'host_type_aliases'
 
 
 class Ohi(object):
+    '''
+        Class for managing openshift host inventory
+    '''
     def __init__(self):
         self.host_type_aliases = {}
         self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))
@@ -35,6 +40,10 @@ class Ohi(object):
         self.aws = awsutil.AwsUtil(self.host_type_aliases)
 
     def run(self):
+        '''
+            Call into awsutil and retrieve the desired hosts and environments
+        '''
+
         if self.args.list_host_types:
             self.aws.print_host_types()
             return 0
@@ -43,18 +52,24 @@ class Ohi(object):
         if self.args.host_type is not None and \
            self.args.env is not None:
             # Both env and host-type specified
-            hosts = self.aws.get_host_list(host_type=self.args.host_type, \
-                                           envs=self.args.env)
+            hosts = self.aws.get_host_list(host_type=self.args.host_type,
+                                           envs=self.args.env,
+                                           version=self.args.openshift_version,
+                                           cached=self.args.cache_only)
 
         if self.args.host_type is None and \
            self.args.env is not None:
             # Only env specified
-            hosts = self.aws.get_host_list(envs=self.args.env)
+            hosts = self.aws.get_host_list(envs=self.args.env,
+                                           version=self.args.openshift_version,
+                                           cached=self.args.cache_only)
 
         if self.args.host_type is not None and \
            self.args.env is None:
             # Only host-type specified
-            hosts = self.aws.get_host_list(host_type=self.args.host_type)
+            hosts = self.aws.get_host_list(host_type=self.args.host_type,
+                                           version=self.args.openshift_version,
+                                           cached=self.args.cache_only)
 
         if hosts is None:
             # We weren't able to determine what they wanted to do
@@ -69,6 +84,9 @@ class Ohi(object):
         return 0
 
     def parse_config_file(self):
+        '''
+            Parse the config file for ohi
+        '''
         if os.path.isfile(self.config_path):
             config = ConfigParser.ConfigParser()
             config.read(self.config_path)
@@ -85,23 +103,27 @@ class Ohi(object):
 
         parser = argparse.ArgumentParser(description='OpenShift Host Inventory')
 
-        parser.add_argument('--list-host-types', default=False, action='store_true',
-                       help='List all of the host types')
+        parser.add_argument('--list-host-types', default=False, action='store_true', help='List all of the host types')
 
-        parser.add_argument('-e', '--env', action="store",
-                       help="Which environment to use")
+        parser.add_argument('-e', '--env', action="store", help="Which environment to use")
 
-        parser.add_argument('-t', '--host-type', action="store",
-                       help="Which host type to use")
+        parser.add_argument('-t', '--host-type', action="store", help="Which host type to use")
 
-        parser.add_argument('-l', '--user', action='store', default=None,
-                               help='username')
+        parser.add_argument('-l', '--user', action='store', default=None, help='username')
 
+        parser.add_argument('-c', '--cache-only', action='store_true', default=False,
+                            help='Retrieve the host inventory by cache only. Default is false.')
 
-        self.args = parser.parse_args()
+        parser.add_argument('-o', '--openshift-version', action='store', default='2',
+                            help='Specify the openshift version. Default is 2')
 
 
-if __name__ == '__main__':
+        self.args = parser.parse_args()
+
+def main():
+    '''
+    Ohi will do its work here
+    '''
     if len(sys.argv) == 1:
         print "\nError: No options given. Use --help to see the available options\n"
         sys.exit(0)
@@ -110,5 +132,9 @@ if __name__ == '__main__':
         ohi = Ohi()
         exitcode = ohi.run()
         sys.exit(exitcode)
-    except ArgumentError as e:
-        print "\nError: %s\n" % e.message
+    except ArgumentError as err:
+        print "\nError: %s\n" % err.message
+
+if __name__ == '__main__':
+    main()
+
diff --git a/bin/openshift_ansible/awsutil.py b/bin/openshift_ansible/awsutil.py
index ba32b4dbd..1ea2f914c 100644
--- a/bin/openshift_ansible/awsutil.py
+++ b/bin/openshift_ansible/awsutil.py
@@ -46,14 +46,17 @@ class AwsUtil(object):
                 self.alias_lookup[value] = key
 
     @staticmethod
-    def get_inventory(args=None):
+    def get_inventory(args=None, cached=False):
         """Calls the inventory script and returns a dictionary containing the inventory."
 
         Keyword arguments:
         args -- optional arguments to pass to the inventory script
         """
         minv = multi_inventory.MultiInventory(args)
-        minv.run()
+        if cached:
+            minv.get_inventory_from_cache()
+        else:
+            minv.run()
         return minv.result
 
     def get_environments(self):
@@ -168,11 +171,12 @@ class AwsUtil(object):
         host_type = self.resolve_host_type(host_type)
         return "tag_env-host-type_%s-%s" % (env, host_type)
 
-    def get_host_list(self, host_type=None, envs=None):
+    def get_host_list(self, host_type=None, envs=None, version=None, cached=False):
         """Get the list of hosts from the inventory using host-type and environment
         """
+        retval = set([])
         envs = envs or []
-        inv = self.get_inventory()
+        inv = self.get_inventory(cached=cached)
 
         # We prefer to deal with a list of environments
         if issubclass(type(envs), basestring):
@@ -183,29 +187,25 @@ class AwsUtil(object):
 
         if host_type and envs:
             # Both host type and environment were specified
-            retval = []
             for env in envs:
-                env_host_type_tag = self.gen_env_host_type_tag(host_type, env)
-                if env_host_type_tag in inv.keys():
-                    retval += inv[env_host_type_tag]
-            return set(retval)
+                retval.update(inv.get('tag_environment_%s' % env, []))
+            retval.intersection_update(inv.get(self.gen_host_type_tag(host_type), []))
 
-        if envs and not host_type:
+        elif envs and not host_type:
             # Just environment was specified
-            retval = []
             for env in envs:
                 env_tag = AwsUtil.gen_env_tag(env)
                 if env_tag in inv.keys():
-                    retval += inv[env_tag]
-            return set(retval)
+                    retval.update(inv.get(env_tag, []))
 
-        if host_type and not envs:
+        elif host_type and not envs:
             # Just host-type was specified
-            retval = []
             host_type_tag = self.gen_host_type_tag(host_type)
             if host_type_tag in inv.keys():
-                retval = inv[host_type_tag]
-            return set(retval)
+                retval.update(inv.get(host_type_tag, []))
+
+        # If version is specified then return only hosts in that version
+        if version:
+            retval.intersection_update(inv.get('oo_version_%s' % version, []))
 
-        # We should never reach here!
-        raise ArgumentError("Invalid combination of parameters")
+        return retval