Merge pull request #151 from twiest/pr

added the ability to run opssh and ohi on all hosts in an environment, as well as all hosts of the same host-type regardless of environment

3 files changed, 88 insertions, 24 deletions

diff --git a/bin/ohi b/bin/ohi
index 06a375cdb..408961ee4 100755
--- a/bin/ohi
+++ b/bin/ohi
@@ -12,13 +12,12 @@ import subprocess
 import ConfigParser
 
 from openshift_ansible import awsutil
+from openshift_ansible.awsutil import ArgumentError
 
 CONFIG_MAIN_SECTION = 'main'
 CONFIG_HOST_TYPE_ALIAS_SECTION = 'host_type_aliases'
 CONFIG_INVENTORY_OPTION = 'inventory'
 
-class ArgumentMismatchError(ValueError): pass
-
 class Ohi(object):
     def __init__(self):
         self.inventory = None
@@ -40,15 +39,30 @@ class Ohi(object):
             self.aws.print_host_types()
             return 0
 
-        if self.args.env and \
-           self.args.host_type:
-            hosts = self.aws.get_host_list(self.args.host_type, self.args.env)
-            for host in hosts:
-                print host
-            return 0
+        hosts = None
+        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, \
+                                           env=self.args.env)
+
+        if self.args.host_type is None and \
+           self.args.env is not None:
+            # Only env specified
+            hosts = self.aws.get_host_list(env=self.args.env)
+
+        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)
+
+        if hosts is None:
+            # We weren't able to determine what they wanted to do
+            raise ArgumentError("Invalid combination of arguments")
 
-        # If it makes it here, we weren't able to determine what they wanted to do
-        raise ArgumentMismatchError("Invalid combination of arguments")
+        for host in hosts:
+            print host
+        return 0
 
     def parse_config_file(self):
         if os.path.isfile(self.config_path):
@@ -92,5 +106,5 @@ if __name__ == '__main__':
         ohi = Ohi()
         exitcode = ohi.run()
         sys.exit(exitcode)
-    except ArgumentMismatchError as e:
+    except ArgumentError as e:
         print "\nError: %s\n" % e.message
diff --git a/bin/openshift_ansible/awsutil.py b/bin/openshift_ansible/awsutil.py
index 8fef0a24f..65b269930 100644
--- a/bin/openshift_ansible/awsutil.py
+++ b/bin/openshift_ansible/awsutil.py
@@ -5,6 +5,10 @@ import os
 import json
 import re
 
+class ArgumentError(Exception):
+    def __init__(self, message):
+        self.message = message
+
 class AwsUtil(object):
     def __init__(self, inventory_path=None, host_type_aliases={}):
         self.host_type_aliases = host_type_aliases
@@ -128,15 +132,45 @@ class AwsUtil(object):
             return self.alias_lookup[host_type]
         return host_type
 
+    def gen_env_tag(self, env):
+        """Generate the environment tag
+        """
+        return "tag_environment_%s" % env
+
+    def gen_host_type_tag(self, host_type):
+        """Generate the host type tag
+        """
+        host_type = self.resolve_host_type(host_type)
+        return "tag_host-type_%s" % host_type
+
     def gen_env_host_type_tag(self, host_type, env):
         """Generate the environment host type tag
         """
         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, env):
+    def get_host_list(self, host_type=None, env=None):
         """Get the list of hosts from the inventory using host-type and environment
         """
         inv = self.get_inventory()
-        host_type_tag = self.gen_env_host_type_tag(host_type, env)
-        return inv[host_type_tag]
+
+        if host_type is not None and \
+           env is not None:
+            # Both host type and environment were specified
+            env_host_type_tag = self.gen_env_host_type_tag(host_type, env)
+            return inv[env_host_type_tag]
+
+        if host_type is None and \
+           env is not None:
+            # Just environment was specified
+            host_type_tag = self.gen_env_tag(env)
+            return inv[host_type_tag]
+
+        if host_type is not None and \
+           env is None:
+            # Just host-type was specified
+            host_type_tag = self.gen_host_type_tag(host_type)
+            return inv[host_type_tag]
+
+        # We should never reach here!
+        raise ArgumentError("Invalid combination of parameters")
diff --git a/bin/opssh b/bin/opssh
index 453da65b4..5fb447318 100755
--- a/bin/opssh
+++ b/bin/opssh
@@ -12,6 +12,7 @@ import subprocess
 import ConfigParser
 
 from openshift_ansible import awsutil
+from openshift_ansible.awsutil import ArgumentError
 
 DEFAULT_PSSH_PAR = 200
 PSSH = '/usr/bin/pssh'
@@ -19,8 +20,6 @@ CONFIG_MAIN_SECTION = 'main'
 CONFIG_HOST_TYPE_ALIAS_SECTION = 'host_type_aliases'
 CONFIG_INVENTORY_OPTION = 'inventory'
 
-class ArgumentMismatchError(ValueError): pass
-
 class Opssh(object):
     def __init__(self):
         self.inventory = None
@@ -42,13 +41,30 @@ class Opssh(object):
             self.aws.print_host_types()
             return 0
 
-        if self.args.env and \
-           self.args.host_type and \
-           self.args.command:
-            return self.run_pssh()
+        hosts = None
+        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, \
+                                           env=self.args.env)
+
+        if self.args.host_type is None and \
+           self.args.env is not None:
+            # Only env specified
+            hosts = self.aws.get_host_list(env=self.args.env)
+
+        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)
+
+        if hosts is None:
+            # We weren't able to determine what they wanted to do
+            raise ArgumentError("Invalid combination of arguments")
 
-        # If it makes it here, we weren't able to determine what they wanted to do
-        raise ArgumentMismatchError("Invalid combination of arguments")
+        for host in hosts:
+            print host
+        return 0
 
     def run_pssh(self):
         """Actually run the pssh command based off of the supplied options
@@ -109,7 +125,7 @@ class Opssh(object):
         parser.add_argument('-e', '--env', action="store",
                        help="Which environment to use")
 
-        parser.add_argument('-t', '--host-type', action="store",
+        parser.add_argument('-t', '--host-type', action="store", default=None,
                        help="Which host type to use")
 
         parser.add_argument('-c', '--command', action='store',
@@ -142,5 +158,5 @@ if __name__ == '__main__':
         opssh = Opssh()
         exitcode = opssh.run()
         sys.exit(exitcode)
-    except ArgumentMismatchError as e:
+    except ArgumentError as e:
         print "\nError: %s\n" % e.message