Skip to content

Commit

Permalink
Removed click from core and better exceptions
Browse files Browse the repository at this point in the history
  • Loading branch information
@wnkz
wnkz committed Mar 27, 2017
1 parent 7e80e35 commit 337746c
Show file tree
Hide file tree
Showing 3 changed files with 131 additions and 71 deletions.
36 changes: 29 additions & 7 deletions shipami/cli.py
View file
Expand Up @@ -81,7 +81,10 @@ def f(_):
return f

now = datetime.datetime.utcnow()
images = shipami.list()
try:
images = shipami.list()
except RuntimeError as e:
raise click.ClickException(str(e))

for k, v in filter_:
attr = headers_mapping.get(filters_mapping.get(k))
Expand Down Expand Up @@ -123,7 +126,10 @@ def f(_):
@click.argument('image-id')
@click.pass_obj
def show(shipami, image_id):
image = shipami.show(image_id)
try:
image = shipami.show(image_id)
except RuntimeError as e:
raise click.ClickException(str(e))

click.echo('id:\t{}'.format(image.get('ImageId')))
click.echo('name:\t{}'.format(image.get('Name')))
Expand Down Expand Up @@ -171,7 +177,11 @@ def show(shipami, image_id):
@click.option('--wait/--no-wait', default=False)
@click.pass_obj
def copy(shipami, **kwargs):
image_id = shipami.copy(kwargs.pop('image_id'), **kwargs)
try:
image_id = shipami.copy(kwargs.pop('image_id'), **kwargs)
except RuntimeError as e:
raise click.ClickException(str(e))

click.echo(image_id)


Expand All @@ -187,7 +197,11 @@ def copy(shipami, **kwargs):
@click.option('--wait/--no-wait', default=False)
@click.pass_obj
def release(shipami, **kwargs):
image_id = shipami.release(kwargs.pop('image_id'), kwargs.pop('release'), **kwargs)
try:
image_id = shipami.release(kwargs.pop('image_id'), kwargs.pop('release'), **kwargs)
except RuntimeError as e:
raise click.ClickException(str(e))

click.echo(image_id)


Expand All @@ -197,13 +211,21 @@ def release(shipami, **kwargs):
@click.option('--remove', is_flag=True, default=False)
@click.pass_obj
def share(shipami, **kwargs):
shipami.share(kwargs.pop('image_id'), **kwargs)
try:
shipami.share(kwargs.pop('image_id'), **kwargs)
except RuntimeError as e:
raise click.ClickException(str(e))


@cli.command()
@click.argument('image-id', nargs=-1)
@click.option('--force', '-f', is_flag=True, default=False)
@click.pass_obj
def delete(shipami, image_id, force):
for deleted in shipami.delete(image_id, force):
click.echo(deleted)
try:
deleted = shipami.delete(image_id, force)
except RuntimeError as e:
raise click.ClickException(str(e))

for d in deleted:
click.echo(d)
158 changes: 98 additions & 60 deletions shipami/core.py
View file
@@ -1,6 +1,5 @@
import json
import logging
import click
import boto3
import botocore
import time
Expand Down Expand Up @@ -33,25 +32,29 @@ def validate_ami_name(self, name, clean=False):
allowed = ['(', ')', '[', ']', ' ', '.', '/', '-', '\'', '@', '_']

if len(name) < 3 or len(name) > 128:
raise click.BadParameter('AMI Name must be 3-128 long (got: "{}")'.format(name))
raise RuntimeError('AMI Name must be 3-128 long (got: "{}")'.format(name))

if clean:
return ''.join(map(lambda _: _ if _.isalnum() or _ in allowed else '-', name))
return name

def list(self):
ec2 = self.__get_session().client('ec2')

r = ec2.describe_images(
Owners=[
'self'
]
)

result = []
copied_keys = ['ImageId', 'Name', 'State', 'CreationDate']
ec2 = self.__get_session().client('ec2')

try:
r = ec2.describe_images(
Owners=[
'self'
]
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

images = r['Images']
images = r.get('Images', [])
for image in images:
i = {}
for key in copied_keys:
Expand Down Expand Up @@ -80,14 +83,16 @@ def show(self, image_id):
]
)
except botocore.exceptions.ClientError as e:
logger.error(e)
raise click.Abort
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

try:
result_image = r.get('Images', [])[0]
except IndexError as e:
logger.error(e)
raise click.Abort
except IndexError:
message = 'Something went wrong'
logger.error(message)
raise RuntimeError(message)

result_image['Shares'] = self.__get_image_permissions(image)
for share in result_image['Shares']:
Expand Down Expand Up @@ -130,8 +135,9 @@ def delete(self, image_ids, force=False):
managed = self.__is_managed(image)
if not managed:
if not force:
logger.error('AMI [{}] is not managed by shipami'.format(image.id))
raise click.Abort
message = 'AMI [{}] is not managed by shipami'.format(image.id)
logger.error(message)
raise RuntimeError(message)
else:
copied_from = self.__get_tag(image, 'shipami:copied_from')
if copied_from:
Expand All @@ -146,8 +152,9 @@ def delete(self, image_ids, force=False):
logger.debug('deleting {}'.format(snapshot.id))
snapshot.delete()
except botocore.exceptions.ClientError as e:
logger.error(e)
raise click.Abort
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

if managed and copied_from:
self.__remove_copied_to(from_image, remove_copied_to)
Expand All @@ -164,20 +171,22 @@ def __copy_image(self, src_image, name=None, name_suffix=None, description=None,
description = description or src_image.description
src_region = src_image.meta.client.meta.region_name
except botocore.exceptions.ClientError as e:
logger.error(e)
raise click.Abort
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

logger.debug('copying image {} from {} to {}'.format(src_image.id, src_region, self._region))
try:
logger.debug('copying image {} from {} to {}'.format(src_image.id, src_region, self._region))
r = self.__get_session().client('ec2').copy_image(
SourceRegion=src_region,
SourceImageId=src_image.id,
Name=name,
Description=description
)
except botocore.exceptions.ClientError as e:
logger.error(e)
raise click.Abort
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

dst_image = self.__get_session().resource('ec2').Image(r['ImageId'])
self.__append_tag(src_image, 'shipami:copied_to', '{}:{}'.format(self._region, dst_image.id))
Expand Down Expand Up @@ -208,16 +217,25 @@ def __copy_tags(self, src_image, dst_image, copy_to_snapshots=False):
snapshot.create_tags(Tags=src_image.tags)

def __share_modify_attribute(self, obj, attribute, operation, account_id):
obj.modify_attribute(
Attribute=attribute,
OperationType=operation,
UserIds=[
account_id
]
)
try:
obj.modify_attribute(
Attribute=attribute,
OperationType=operation,
UserIds=[
account_id
]
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

def __get_copied_from_image(self, copied_from):
region, image_id = copied_from.split(':')
try:
region, image_id = copied_from.split(':')
except ValueError as e:
logger.error(e)
raise RuntimeError(e)
image = self.__get_session(region).resource('ec2').Image(image_id)
return image

Expand Down Expand Up @@ -268,7 +286,7 @@ def __append_tag(self, obj, key, value):
def __get_tag(self, obj, key):
try:
tags = obj.tags or []
except AttributeError:
except (botocore.exceptions.ClientError, AttributeError):
try:
tags = obj.get('Tags', [])
except AttributeError:
Expand All @@ -280,45 +298,65 @@ def __get_tag(self, obj, key):
return None

def __set_tag(self, obj, key, value):
obj.create_tags(
Tags=[
{
'Key': key,
'Value': value
}
]
)
try:
obj.create_tags(
Tags=[
{
'Key': key,
'Value': value
}
]
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

def __delete_tag(self, obj, key):
obj.meta.client.delete_tags(
Resources=[
obj.id
],
Tags=[
{
'Key': key
}
]
)
try:
obj.meta.client.delete_tags(
Resources=[
obj.id
],
Tags=[
{
'Key': key
}
]
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

def __is_managed(self, image):
if self.__get_tag(image, 'shipami:managed') == 'True':
return True
return False

def __get_image_permissions(self, image):
r = image.describe_attribute(
Attribute='launchPermission'
)
try:
r = image.describe_attribute(
Attribute='launchPermission'
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

return r['LaunchPermissions']
return r.get('LaunchPermissions', [])

def __get_snapshot_permissions(self, snapshot):
r = snapshot.describe_attribute(
Attribute='createVolumePermission'
)
try:
r = snapshot.describe_attribute(
Attribute='createVolumePermission'
)
except botocore.exceptions.ClientError as e:
message = e.response['Error']['Message']
logger.error(message)
raise RuntimeError(message)

return r['CreateVolumePermissions']
return r.get('CreateVolumePermissions', [])

def __is_ami_shared(self, image, account_id=None):
account_id = account_id or self.MARKETPLACE_ACCOUNT_ID
Expand Down
8 changes: 4 additions & 4 deletions tests/test_cli.py
View file
Expand Up @@ -147,7 +147,7 @@ def test_show_inexistant_id(self):
r = runner.invoke(shipami, ['show', 'ami-42424242'])

assert r.exit_code == 1
assert 'Aborted!' in r.output
assert 'Error: The image id \'[ami-42424242]\' does not exist' in r.output

def test_copy(self, ec2, base_image):
image_number = len(ec2.meta.client.describe_images()['Images'])
Expand Down Expand Up @@ -202,7 +202,7 @@ def test_copy_invalid_name(self, base_image):

r = runner.invoke(shipami, ['copy', base_image.id, '--name', NAME])

assert 'Invalid value' in r.output
assert 'Error:' in r.output

def test_copy_cleanup_name(self, ec2, base_image):
NAME = 'foo#bar'
Expand All @@ -219,7 +219,7 @@ def test_copy_inexistant_id(self, ec2, base_image):
r = runner.invoke(shipami, ['copy', 'ami-42424242'])

assert r.exit_code == 1
assert 'Aborted!' in r.output
assert 'Error: The image id \'[ami-42424242]\' does not exist' in r.output

def test_release(self, ec2, base_image):
RELEASE = '1.0.0'
Expand Down Expand Up @@ -266,7 +266,7 @@ def test_delete_error(self, base_image):
r = runner.invoke(shipami, ['delete', base_image_id])

assert r.exit_code == 1
assert 'Aborted!' in r.output
assert 'Error: AMI [{}] is not managed by shipami'.format(base_image_id) in r.output

def test_delete_force(self, ec2, base_image):
base_image_id = base_image.id
Expand Down

0 comments on commit 337746c

Please sign in to comment.