Merge pull request #248 from spenczar/88692b946a1b4f4f20575e308958b1c56f6ede6c

Added AddInstanceGroup and ModifyInstanceGroup to boto.emr

3 files changed, 101 insertions, 0 deletions

diff --git a/boto/emr/connection.py b/boto/emr/connection.py
index 70ddf824..0d3cb8b3 100644
--- a/boto/emr/connection.py
+++ b/boto/emr/connection.py
@@ -145,6 +145,49 @@ class EmrConnection(AWSQueryConnection):
         return self.get_object(
             'AddJobFlowSteps', params, RunJobFlowResponse, verb='POST')
 
+	def add_instance_groups(self, jobflow_id, instance_groups):
+		"""
+		Adds instance groups to a running cluster.
+		
+		:type jobflow_id: str
+		:param jobflow_id: The id of the jobflow which will take the new instance groups
+		:type instance_groups: list(boto.emr.InstanceGroup)
+		:param instance_groups: A list of instance groups to add to the job
+		"""
+		if type(instance_groups) != types.ListType:
+			instance_groups = [instance_groups]
+		params = {}
+		params['JobFlowId'] = jobflow_id
+		
+		instance_group_args = [self._build_instance_group_args(ig) for ig in instance_groups]
+		params.update(self._build_instance_group_list(instance_group_args))
+		
+		return self.get_object('AddInstanceGroups', params, AddInstanceGroupsResponse, verb='POST')
+			
+	def modify_instance_groups(self, instance_group_ids, new_sizes):
+		"""
+		Modify the number of nodes and configuration settings in an instance group.
+				
+		:type instance_group_ids: list(str)
+		:param instance_group_ids: A list of the ID's of the instance groups to be modified
+		:type new_sizes: list(int)
+		:param new_sizes: A list of the new sizes for each instance group
+		"""
+		if type(instance_group_ids) != types.ListType:
+			instance_group_ids = [instance_group_ids]
+		if type(new_sizes) != types.ListType:
+			new_sizes = [new_sizes]
+			
+		instance_groups = zip(instance_group_ids, new_sizes)
+				
+		for k, ig in enumerate(instance_groups):
+			#could be wrong - the example amazon gives uses InstanceRequestCount, 
+			#while the api documentation says InstanceCount
+			params['InstanceGroups.member.%s.InstanceGroupId' % k+1 ] = ig[1][0] 
+			params['InstanceGroups.member.%s.InstanceCount' % k+1 ] = ig[1][1]
+			
+		return self.get_object('ModifyInstanceGroups', params, ModifyInstanceGroupsResponse, verb='POST')
+
     def run_jobflow(self, name, log_uri, ec2_keyname=None, availability_zone=None,
                     master_instance_type='m1.small',
                     slave_instance_type='m1.small', num_instances=1,
@@ -284,3 +327,22 @@ class EmrConnection(AWSQueryConnection):
 
         return params
 
+    def _build_instance_group_args(self, instance_group):
+		params = {
+			'InstanceCount' : instance_group.num_instances,
+			'InstanceRole' : instance_group.role,
+			'InstanceType' : instance_group.type,
+			'Name' : instance_group.name,
+			'Market' : instance_group.market
+		}
+		return params
+		
+    def _build_instance_group_list(self, instance_groups):
+        if type(instance_groups) != types.ListType:
+            instance_groups = [instance_groups]
+
+        params = {}
+        for i, instance_group in enumerate(instance_groups):
+            for key, value in instance_group.iteritems():
+                params['InstanceGroups.member.%s.%s' % (i+1, key)] = value
+        return params
diff --git a/boto/emr/emrobject.py b/boto/emr/emrobject.py
index e5887753..3430b987 100644
--- a/boto/emr/emrobject.py
+++ b/boto/emr/emrobject.py
@@ -45,6 +45,12 @@ class EmrObject(object):
 class RunJobFlowResponse(EmrObject):
     Fields = set(['JobFlowId'])
 
+class AddInstanceGroupsResponse(EmrObject):
+    Fields = set(['InstanceGroupIds', 'JobFlowId'])
+    
+class ModifyInstanceGroupsResponse(EmrObject):
+    Fields = set(['RequestId'])
+    
 
 class Arg(EmrObject):
     def __init__(self, connection=None):
diff --git a/boto/emr/instance_group.py b/boto/emr/instance_group.py
new file mode 100644
index 00000000..baa11cea
--- /dev/null
+++ b/boto/emr/instance_group.py
@@ -0,0 +1,33 @@
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish, dis-
+# tribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the fol-
+# lowing conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
+# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+
+class InstanceGroup(object):
+	def __init__(self, num_instances, role, type, market, name):
+		self.num_instances = num_instances
+		self.role = role
+		self.type = type
+		self.market = market
+		self.name = name
+	
+	def __repr__(self):
+		return '%s.%s(name=%r, num_instances=%r, role=%r, type=%r, market = %r)' % (
+            self.__class__.__module__, self.__class__.__name__,
+            self.name, self.num_instances, self.role, self.type, self.market)