diff --git a/senza/aws.py b/senza/aws.py index a2be35de..037ad055 100644 --- a/senza/aws.py +++ b/senza/aws.py @@ -301,10 +301,13 @@ def get_stacks(stack_refs: list, # stack names that were already yielded to avoid yielding old deleted # stacks whose name was reused stacks_yielded = set() + output_stacks = [] for stack in stacks: if stack['StackName'] not in stacks_yielded or not unique_only: stacks_yielded.add(stack['StackName']) - yield SenzaStackSummary(stack) + output_stacks.append(SenzaStackSummary(stack)) + + return output_stacks def matches_any(cf_stack_name: str, stack_refs: list): @@ -461,7 +464,7 @@ def all_stacks_in_final_state(related_stacks_refs: list, region: str, timeout: O while not all_in_final_state and wait_timeout > datetime.datetime.utcnow(): # assume all stacks are ready all_in_final_state = True - related_stacks = list(get_stacks(related_stacks_refs, region)) + related_stacks = get_stacks(related_stacks_refs, region) if not related_stacks: error("Stack not found!") diff --git a/senza/cli.py b/senza/cli.py index 12026792..ec237240 100755 --- a/senza/cli.py +++ b/senza/cli.py @@ -25,6 +25,7 @@ fatal_error, info, ok) from clickclick.console import print_table +from spotinst.components import elastigroup_api from .arguments import (GLOBAL_OPTIONS, json_output_option, output_option, parameter_file_option, region_option, stacktrace_visible_option, watch_option, @@ -118,6 +119,12 @@ SENZA_KMS_PREFIX = 'senza:kms:' +ELASTIGROUP_TYPE = 'Custom::elastigroup' + +AUTO_SCALING_GROUP_TYPE = 'AWS::AutoScaling::AutoScalingGroup' + +VALID_AUTO_SCALING_GROUPS = [AUTO_SCALING_GROUP_TYPE, ELASTIGROUP_TYPE] + def filter_output_columns(output_columns, filter_columns): """ @@ -1426,6 +1433,23 @@ def dump(stack_ref, region, output): print_json(cfjson, output) +def get_auto_scaling_groups_and_elasti_groups(stacks, region): + ''' + Returns data for both AWS Auto Scaling Groups and ElastiGroups + + Note: This method will eventually replace get_auto_scaling_groups when the remaining commands support ElastiGroups + ''' + cf = BotoClientProxy('cloudformation', region) + for stack in stacks: + resources = cf.describe_stack_resources(StackName=stack.StackName)['StackResources'] + + for resource in resources: + if resource['ResourceType'] in VALID_AUTO_SCALING_GROUPS: + yield {'type': resource['ResourceType'], + 'resource_id': resource['PhysicalResourceId'], + 'stack_name': resource['StackName']} + + def get_auto_scaling_groups(stack_refs, region): cf = BotoClientProxy('cloudformation', region) for stack in get_stacks(stack_refs, region): @@ -1523,27 +1547,67 @@ def scale(stack_ref, region, desired_capacity, force): asg = BotoClientProxy('autoscaling', region) - stack_count = len(list(get_stacks(stack_refs, region))) + stacks = get_stacks(stack_refs, region) + stack_count = len(stacks) if not force and stack_count > 1: confirm_str = 'Number of stacks to be scaled - {}. Do you want to continue?'.format(stack_count) click.confirm(confirm_str, abort=True) - for asg_name in get_auto_scaling_groups(stack_refs, region): - group = get_auto_scaling_group(asg, asg_name) - current_capacity = group['DesiredCapacity'] - with Action('Scaling {} from {} to {} instances..'.format( - asg_name, current_capacity, desired_capacity)) as act: - if current_capacity == desired_capacity: - act.ok('NO CHANGES') - else: - kwargs = {} - if desired_capacity < group['MinSize']: - kwargs['MinSize'] = desired_capacity - if desired_capacity > group['MaxSize']: - kwargs['MaxSize'] = desired_capacity - asg.update_auto_scaling_group(AutoScalingGroupName=asg_name, - DesiredCapacity=desired_capacity, - **kwargs) + for group in get_auto_scaling_groups_and_elasti_groups(stacks, region): + if group['type'] == AUTO_SCALING_GROUP_TYPE: + scale_auto_scaling_group(asg, group['resource_id'], desired_capacity) + elif group['type'] == ELASTIGROUP_TYPE: + scale_elastigroup(group['resource_id'], group['stack_name'], desired_capacity, region) + + +def scale_elastigroup(elastigroup_id, stack_name, desired_capacity, region): + ''' + Commands to scale an ElastiGroup + ''' + cf = boto3.client('cloudformation', region) + template = cf.get_template(StackName=stack_name)['TemplateBody'] + + spotinst_token = template['Mappings']['Senza']['Info']['SpotinstAccessToken'] + spotinst_account_id = template['Resources']['AppServerConfig']['Properties']['accountId'] + + group = elastigroup_api.get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token) + capacity = group['capacity'] + + with Action('Scaling ElastiGroup {} (ID: {}) from {} to {} instances..'.format( + stack_name, elastigroup_id, capacity['target'], desired_capacity)) as act: + if capacity['target'] == desired_capacity: + act.ok('NO CHANGES') + else: + minimum = desired_capacity if capacity['minimum'] > desired_capacity else capacity['minimum'] + maximum = desired_capacity if capacity['maximum'] < desired_capacity else capacity['maximum'] + + elastigroup_api.update_capacity(minimum, + maximum, + desired_capacity, + elastigroup_id, + spotinst_account_id, + spotinst_token) + + +def scale_auto_scaling_group(asg, asg_name, desired_capacity): + ''' + Commands to scale an AWS Auto Scaling Group + ''' + group = get_auto_scaling_group(asg, asg_name) + current_capacity = group['DesiredCapacity'] + with Action('Scaling {} from {} to {} instances..'.format( + asg_name, current_capacity, desired_capacity)) as act: + if current_capacity == desired_capacity: + act.ok('NO CHANGES') + else: + kwargs = {} + if desired_capacity < group['MinSize']: + kwargs['MinSize'] = desired_capacity + if desired_capacity > group['MaxSize']: + kwargs['MaxSize'] = desired_capacity + asg.update_auto_scaling_group(AutoScalingGroupName=asg_name, + DesiredCapacity=desired_capacity, + **kwargs) def failure_event(event: dict): @@ -1586,7 +1650,7 @@ def wait(stack_ref, region, deletion, timeout, interval): stacks_in_progress = set() successful_actions = set() - stacks_found = list(get_stacks(stack_refs, region, all=True, unique_only=True)) + stacks_found = get_stacks(stack_refs, region, all=True, unique_only=True) if not stacks_found: raise click.UsageError('No matching stack for "{}" found'.format(' '.join(stack_ref))) diff --git a/spotinst/components/elastigroup_api.py b/spotinst/components/elastigroup_api.py new file mode 100644 index 00000000..063a251d --- /dev/null +++ b/spotinst/components/elastigroup_api.py @@ -0,0 +1,60 @@ +''' +Wrapper methods for ElastiGroup's API +''' +import requests +import json + + +SPOTINST_API_URL = 'https://api.spotinst.io' + + +def update_capacity(minimum, maximum, target, elastigroup_id, spotinst_account_id, spotinst_token): + ''' + Updates the capacity (number of instances) for an ElastiGroup by calling the SpotInst API. + Returns the updated description of the ElastiGroup as a dict. + Exceptions will be thrown for HTTP errors. + + For more details see https://api.spotinst.com/elastigroup/amazon-web-services/update/ + ''' + headers = { + "Authorization": "Bearer {}".format(spotinst_token), + "Content-Type": "application/json" + } + + new_capacity = { + 'target': target, + 'minimum': minimum, + 'maximum': maximum + } + + body = {'group': {'capacity': new_capacity}} + response = requests.put( + '{}/aws/ec2/group/{}?accountId={}'.format(SPOTINST_API_URL, elastigroup_id, spotinst_account_id), + headers=headers, timeout=10, data=json.dumps(body)) + response.raise_for_status() + data = response.json() + groups = data.get("response", {}).get("items", []) + + return groups[0] + + +def get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token): + ''' + Returns the description of an ElastiGroup as a dict. + Exceptions will be thrown for HTTP errors. + + For more details see https://api.spotinst.com/elastigroup/amazon-web-services/list-group/ + ''' + headers = { + "Authorization": "Bearer {}".format(spotinst_token), + "Content-Type": "application/json" + } + + response = requests.get( + 'https://api.spotinst.io/aws/ec2/group/{}?accountId={}'.format(elastigroup_id, spotinst_account_id), + headers=headers, timeout=5) + response.raise_for_status() + data = response.json() + groups = data.get("response", {}).get("items", []) + + return groups[0] diff --git a/tests/test_cli.py b/tests/test_cli.py index e4f04943..2120a4e3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -5,7 +5,6 @@ from contextlib import contextmanager from typing import List, Dict -from typing import Optional from unittest.mock import MagicMock, mock_open import botocore.exceptions @@ -1645,7 +1644,8 @@ def test_scale(monkeypatch): 'CreationTime': '2016-06-14'}]} boto3.describe_stack_resources.return_value = {'StackResources': [{'ResourceType': 'AWS::AutoScaling::AutoScalingGroup', - 'PhysicalResourceId': 'myasg'}]} + 'PhysicalResourceId': 'myasg', + 'StackName': 'myapp-1'}]} # NOTE: we are using invalid MinSize (< capacity) here to get one more line covered ;-) group = {'AutoScalingGroupName': 'myasg', 'DesiredCapacity': 1, 'MinSize': 3, 'MaxSize': 1} boto3.describe_auto_scaling_groups.return_value = {'AutoScalingGroups': [group]} @@ -1655,6 +1655,68 @@ def test_scale(monkeypatch): catch_exceptions=False) assert 'Scaling myasg from 1 to 2 instances' in result.output + +def test_scale_elastigroup(monkeypatch): + spotinst_account_id = 'fakeactid' + elastigroup_id = 'myelasti' + boto3 = MagicMock() + boto3.list_stacks.return_value = {'StackSummaries': [{'StackName': 'myapp-1', + 'CreationTime': '2016-06-14'}]} + boto3.describe_stack_resources.return_value = {'StackResources': + [{'ResourceType': 'Custom::elastigroup', + 'PhysicalResourceId': elastigroup_id, + 'StackName': 'myapp-1'}]} + boto3.get_template.return_value = { + 'TemplateBody': { + 'Mappings': { + 'Senza': { + 'Info': { + 'SpotinstAccessToken': 'faketoken' + } + } + }, + 'Resources': { + 'AppServerConfig': { + 'Properties': { + 'accountId': spotinst_account_id + } + } + } + } + } + + group = { + 'capacity': { + 'minimum': 1, + 'maximum': 2, + 'target': 1, + 'unit': 'instance' + } + } + get_elastigroup = MagicMock() + get_elastigroup.return_value = group + + update = { + 'capacity': { + 'minimum': 1, + 'maximum': 3, + 'target': 3, + 'unit': 'instance' + } + } + update_capacity = MagicMock() + update_capacity.return_value = update + + monkeypatch.setattr('boto3.client', MagicMock(return_value=boto3)) + monkeypatch.setattr('spotinst.components.elastigroup_api.get_elastigroup', get_elastigroup) + monkeypatch.setattr('spotinst.components.elastigroup_api.update_capacity', update_capacity) + + runner = CliRunner() + result = runner.invoke(cli, ['scale', 'myapp', '1', '3', '--region=aa-fakeregion-1'], + catch_exceptions=False) + assert 'Scaling ElastiGroup myapp-1 (ID: myelasti) from 1 to 3 instances' in result.output + + def test_scale_with_confirm(monkeypatch): boto3 = MagicMock() boto3.list_stacks.return_value = {'StackSummaries': [ @@ -1683,7 +1745,8 @@ def test_scale_with_force_confirm(monkeypatch): ]} boto3.describe_stack_resources.return_value = {'StackResources': [{'ResourceType': 'AWS::AutoScaling::AutoScalingGroup', - 'PhysicalResourceId': 'myasg'}]} + 'PhysicalResourceId': 'myasg', + 'StackName': 'myapp-1'}]} group = {'AutoScalingGroupName': 'myasg', 'DesiredCapacity': 1, 'MinSize': 3, 'MaxSize': 1} boto3.describe_auto_scaling_groups.return_value = {'AutoScalingGroups': [group]} monkeypatch.setattr('boto3.client', MagicMock(return_value=boto3)) @@ -1693,7 +1756,6 @@ def test_scale_with_force_confirm(monkeypatch): assert 'Scaling myasg from 1 to 2 instances' in result.output - def test_wait(monkeypatch): cf = MagicMock() stack1 = {'StackName': 'test-1', diff --git a/tests/test_elastigroup_api.py b/tests/test_elastigroup_api.py new file mode 100644 index 00000000..39f92a48 --- /dev/null +++ b/tests/test_elastigroup_api.py @@ -0,0 +1,59 @@ +import responses +from spotinst.components.elastigroup_api import update_capacity, get_elastigroup, SPOTINST_API_URL + + +def test_update_capacity(monkeypatch): + update = { + 'response': { + 'items': [{ + 'id': 'sig-xfy', + 'name': 'my-app-1', + 'capacity': { + 'minimum': 1, + 'maximum': 3, + 'target': 3, + 'unit': 'instance' + } + }] + } + } + + elastigroup_id = 'sig-xfy' + spotinst_account_id = 'act-zwk' + spotinst_token = 'fake-token' + with responses.RequestsMock() as rsps: + rsps.add(rsps.PUT, + '{}/aws/ec2/group/{}?accountId={}'.format(SPOTINST_API_URL, elastigroup_id, spotinst_account_id), + status=200, + json=update) + + update_response = update_capacity(1, 3, 3, elastigroup_id, spotinst_account_id, spotinst_token) + assert update_response['id'] == elastigroup_id + assert update_response['name'] == 'my-app-1' + assert update_response['capacity']['minimum'] == 1 + assert update_response['capacity']['maximum'] == 3 + assert update_response['capacity']['target'] == 3 + + +def test_get_elastigroup(monkeypatch): + group = { + 'response': { + 'items': [{ + 'id': 'sig-xfy', + 'name': 'my-app-1', + 'region': 'eu-central-1', + }] + } + } + + elastigroup_id = 'sig-xfy' + spotinst_account_id = 'act-zwk' + spotinst_token = 'fake-token' + with responses.RequestsMock() as rsps: + rsps.add(rsps.GET, 'https://api.spotinst.io/aws/ec2/group/{}?accountId={}'.format(elastigroup_id, spotinst_account_id), + status=200, + json=group) + + group = get_elastigroup(elastigroup_id, spotinst_account_id, spotinst_token) + assert group['id'] == elastigroup_id + assert group['name'] == 'my-app-1'