ohi: added subtype searching

2 files changed, 29 insertions, 1 deletions

diff --git a/bin/ohi b/bin/ohi
index f9e76b783..9c2ce8432 100755
--- a/bin/ohi
+++ b/bin/ohi
@@ -34,6 +34,7 @@ class Ohi(object):
                                         'openshift_ansible', \
                                         'openshift_ansible.conf')
 
+        self.args = None
         self.parse_cli_args()
         self.parse_config_file()
 
@@ -57,6 +58,7 @@ class Ohi(object):
 
         hosts = self.aws.get_host_list(clusters=self.args.cluster,
                                        host_type=self.args.host_type,
+                                       sub_host_type=self.args.sub_host_type,
                                        envs=self.args.env,
                                        version=version,
                                        cached=self.args.cache_only)
@@ -104,6 +106,8 @@ class Ohi(object):
 
         parser.add_argument('-t', '--host-type', action="store", help="Which host type to use")
 
+        parser.add_argument('-s', '--sub-host-type', action="store", help="Which sub host type to use")
+
         parser.add_argument('-l', '--user', action='store', default=None, help='username')
 
         parser.add_argument('--cache-only', action='store_true', default=False,
diff --git a/bin/openshift_ansible/awsutil.py b/bin/openshift_ansible/awsutil.py
index 945e6a20c..d46af3d66 100644
--- a/bin/openshift_ansible/awsutil.py
+++ b/bin/openshift_ansible/awsutil.py
@@ -37,6 +37,7 @@ class AwsUtil(object):
         self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))
 
         self.setup_host_type_alias_lookup()
+        self.alias_lookup = None
 
     def setup_host_type_alias_lookup(self):
         """Sets up the alias to host-type lookup table."""
@@ -101,6 +102,20 @@ class AwsUtil(object):
         host_types.sort()
         return host_types
 
+    def get_sub_host_types(self):
+        """Searches for sub-host-type tags in the inventory and returns all sub-host-types found."""
+        pattern = re.compile(r'^oo_subhosttype_(.*)')
+
+        sub_host_types = []
+        inv = self.get_inventory()
+        for key in inv.keys():
+            matched = pattern.match(key)
+            if matched:
+                sub_host_types.append(matched.group(1))
+
+        sub_host_types.sort()
+        return sub_host_types
+
     def get_security_groups(self):
         """Searches for security_groups in the inventory and returns all SGs found."""
         pattern = re.compile(r'^security_group_(.*)')
@@ -192,9 +207,15 @@ class AwsUtil(object):
             host_type = self.resolve_host_type(host_type)
         return "oo_hosttype_%s" % host_type
 
+    @staticmethod
+    def gen_sub_host_type_tag(sub_host_type):
+        """Generate the host type tag
+        """
+        return "oo_subhosttype_%s" % sub_host_type
+
     # This function uses all of these params to perform a filters on our host inventory.
     # pylint: disable=too-many-arguments
-    def get_host_list(self, clusters=None, host_type=None, envs=None, version=None, cached=False):
+    def get_host_list(self, clusters=None, host_type=None, sub_host_type=None, envs=None, version=None, cached=False):
         """Get the list of hosts from the inventory using host-type and environment
         """
         retval = set([])
@@ -229,6 +250,9 @@ class AwsUtil(object):
         if host_type:
             retval.intersection_update(inv.get(self.gen_host_type_tag(host_type, version), []))
 
+        if sub_host_type:
+            retval.intersection_update(inv.get(self.gen_sub_host_type_tag(sub_host_type), []))
+
         if version != 'all':
             retval.intersection_update(inv.get(AwsUtil.gen_version_tag(version), []))