diff --git a/.gitignore b/.gitignore index a65824a..8b29acb 100644 --- a/.gitignore +++ b/.gitignore @@ -60,3 +60,6 @@ target/ # Jetbrains/PyCharm project files .idea/ + +# vim swap files +.*.sw? diff --git a/aws_lambda/aws_lambda.py b/aws_lambda/aws_lambda.py index 44f37cc..e0b0f03 100755 --- a/aws_lambda/aws_lambda.py +++ b/aws_lambda/aws_lambda.py @@ -33,7 +33,10 @@ log = logging.getLogger(__name__) -def cleanup_old_versions(src, keep_last_versions, config_file='config.yaml'): +def cleanup_old_versions( + src, keep_last_versions, + config_file='config.yaml', profile_name=None, +): """Deletes old deployed versions of the function in AWS Lambda. Won't delete $Latest and any aliased version @@ -48,13 +51,14 @@ def cleanup_old_versions(src, keep_last_versions, config_file='config.yaml'): print("Won't delete all versions. Please do this manually") else: path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) + profile_name = cfg.get('profile') aws_access_key_id = cfg.get('aws_access_key_id') aws_secret_access_key = cfg.get('aws_secret_access_key') client = get_client( - 'lambda', aws_access_key_id, aws_secret_access_key, + 'lambda', profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) @@ -80,7 +84,7 @@ def cleanup_old_versions(src, keep_last_versions, config_file='config.yaml'): def deploy( src, use_requirements=False, local_package=None, - config_file='config.yaml', + config_file='config.yaml', profile_name=None, ): """Deploys a new function to AWS Lambda. @@ -93,7 +97,7 @@ def deploy( """ # Load and parse the config file. path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) # Copy all the pip dependencies required to run your code into a temporary # folder then add the handler file in the root of this directory. @@ -112,7 +116,8 @@ def deploy( def deploy_s3( - src, use_requirements=False, local_package=None, config_file='config.yaml', + src, use_requirements=False, local_package=None, + config_file='config.yaml', profile_name=None, ): """Deploys a new function via AWS S3. @@ -125,7 +130,7 @@ def deploy_s3( """ # Load and parse the config file. path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) # Copy all the pip dependencies required to run your code into a temporary # folder then add the handler file in the root of this directory. @@ -146,7 +151,7 @@ def deploy_s3( def upload( src, use_requirements=False, local_package=None, - config_file='config.yaml', + config_file='config.yaml', profile_name=None, ): """Uploads a new function to AWS S3. @@ -159,7 +164,7 @@ def upload( """ # Load and parse the config file. path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) # Copy all the pip dependencies required to run your code into a temporary # folder then add the handler file in the root of this directory. @@ -174,7 +179,8 @@ def upload( def invoke( - src, event_file='event.json', config_file='config.yaml', + src, event_file='event.json', + config_file='config.yaml', profile_name=None, verbose=False, ): """Simulates a call to your function. @@ -189,7 +195,11 @@ def invoke( """ # Load and parse the config file. path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) + + # Set AWS_PROFILE environment variable based on `--profile` option. + if profile_name: + os.environ['AWS_PROFILE'] = profile_name # Load environment variables from the config file into the actual # environment. @@ -248,7 +258,8 @@ def init(src, minimal=False): def build( - src, use_requirements=False, local_package=None, config_file='config.yaml', + src, use_requirements=False, local_package=None, + config_file='config.yaml', profile_name=None, ): """Builds the file bundle. @@ -261,7 +272,7 @@ def build( """ # Load and parse the config file. path_to_config_file = os.path.join(src, config_file) - cfg = read(path_to_config_file, loader=yaml.load) + cfg = read_cfg(path_to_config_file, profile_name) # Get the absolute path to the output directory and create it if it doesn't # already exist. @@ -439,24 +450,25 @@ def get_role_name(region, account_id, role): return 'arn:{0}:iam::{1}:role/{2}'.format(prefix, account_id, role) -def get_account_id(aws_access_key_id, aws_secret_access_key, region=None): +def get_account_id(profile_name, aws_access_key_id, aws_secret_access_key, region=None): """Query STS for a users' account_id""" client = get_client( - 'sts', aws_access_key_id, aws_secret_access_key, + 'sts', profile_name, aws_access_key_id, aws_secret_access_key, region, ) return client.get_caller_identity().get('Account') -def get_client(client, aws_access_key_id, aws_secret_access_key, region=None): +def get_client(client, profile_name, aws_access_key_id, aws_secret_access_key, region=None): """Shortcut for getting an initialized instance of the boto3 client.""" - return boto3.client( - client, + boto3.setup_default_session( + profile_name=profile_name, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, region_name=region, ) + return boto3.client(client) def create_function(cfg, path_to_zip_file, *use_s3, **s3_file): @@ -464,11 +476,12 @@ def create_function(cfg, path_to_zip_file, *use_s3, **s3_file): print('Creating your new Lambda function') byte_stream = read(path_to_zip_file, binary_file=True) + profile_name = cfg.get('profile') aws_access_key_id = cfg.get('aws_access_key_id') aws_secret_access_key = cfg.get('aws_secret_access_key') account_id = get_account_id( - aws_access_key_id, aws_secret_access_key, cfg.get('region'), + profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) role = get_role_name( cfg.get('region'), account_id, @@ -476,7 +489,7 @@ def create_function(cfg, path_to_zip_file, *use_s3, **s3_file): ) client = get_client( - 'lambda', aws_access_key_id, aws_secret_access_key, + 'lambda', profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) @@ -536,11 +549,12 @@ def update_function(cfg, path_to_zip_file, *use_s3, **s3_file): print('Updating your Lambda function') byte_stream = read(path_to_zip_file, binary_file=True) + profile_name = cfg.get('profile') aws_access_key_id = cfg.get('aws_access_key_id') aws_secret_access_key = cfg.get('aws_secret_access_key') account_id = get_account_id( - aws_access_key_id, aws_secret_access_key, cfg.get('region'), + profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) role = get_role_name( cfg.get('region'), account_id, @@ -548,7 +562,7 @@ def update_function(cfg, path_to_zip_file, *use_s3, **s3_file): ) client = get_client( - 'lambda', aws_access_key_id, aws_secret_access_key, + 'lambda', profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) @@ -603,10 +617,11 @@ def upload_s3(cfg, path_to_zip_file, *use_s3): """Upload a function to AWS S3.""" print('Uploading your new Lambda function') + profile_name = cfg.get('profile') aws_access_key_id = cfg.get('aws_access_key_id') aws_secret_access_key = cfg.get('aws_secret_access_key') client = get_client( - 's3', aws_access_key_id, aws_secret_access_key, + 's3', profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) byte_stream = b'' @@ -641,10 +656,11 @@ def upload_s3(cfg, path_to_zip_file, *use_s3): def function_exists(cfg, function_name): """Check whether a function exists or not""" + profile_name = cfg.get('profile') aws_access_key_id = cfg.get('aws_access_key_id') aws_secret_access_key = cfg.get('aws_secret_access_key') client = get_client( - 'lambda', aws_access_key_id, aws_secret_access_key, + 'lambda', profile_name, aws_access_key_id, aws_secret_access_key, cfg.get('region'), ) @@ -663,3 +679,12 @@ def function_exists(cfg, function_name): f['FunctionName'] for f in functions_resp.get('Functions', []) ]) return function_name in functions + + +def read_cfg(path_to_config_file, profile_name): + cfg = read(path_to_config_file, loader=yaml.load) + if profile_name is not None: + cfg['profile'] = profile_name + elif 'AWS_PROFILE' in os.environ: + cfg['profile'] = os.environ['AWS_PROFILE'] + return cfg diff --git a/scripts/lambda b/scripts/lambda index 3f3f7ae..589eecc 100755 --- a/scripts/lambda +++ b/scripts/lambda @@ -43,6 +43,10 @@ def init(folder, minimal): default='config.yaml', help='Alternate config file.', ) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option( '--use-requirements', default=False, @@ -56,12 +60,13 @@ def init(folder, minimal): help='Install local package as well.', multiple=True, ) -def build(use_requirements, local_package, config_file): +def build(use_requirements, local_package, config_file, profile): aws_lambda.build( CURRENT_DIR, use_requirements=use_requirements, local_package=local_package, config_file=config_file, + profile_name=profile, ) @@ -76,12 +81,17 @@ def build(use_requirements, local_package, config_file): default='config.yaml', help='Alternate config file.', ) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option('--verbose', '-v', is_flag=True) -def invoke(event_file, config_file, verbose): +def invoke(event_file, config_file, profile, verbose): aws_lambda.invoke( CURRENT_DIR, event_file=event_file, config_file=config_file, + profile_name=profile, verbose=verbose, ) @@ -92,6 +102,10 @@ def invoke(event_file, config_file, verbose): default='config.yaml', help='Alternate config file.', ) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option( '--use-requirements', default=False, @@ -105,16 +119,26 @@ def invoke(event_file, config_file, verbose): help='Install local package as well.', multiple=True, ) -def deploy(use_requirements, local_package, config_file): +def deploy(use_requirements, local_package, config_file, profile): aws_lambda.deploy( CURRENT_DIR, - config_file=config_file, use_requirements=use_requirements, local_package=local_package, + config_file=config_file, + profile_name=profile, ) @click.command(help='Upload your lambda to S3.') +@click.option( + '--config-file', + default='config.yaml', + help='Alternate config file.', +) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option( '--use-requirements', default=False, @@ -128,8 +152,14 @@ def deploy(use_requirements, local_package, config_file): help='Install local package as well.', multiple=True, ) -def upload(use_requirements, local_package): - aws_lambda.upload(CURRENT_DIR, use_requirements, local_package) +def upload(use_requirements, local_package, config_file, profile): + aws_lambda.upload( + CURRENT_DIR, + use_requirements=use_requirements, + local_package=local_package, + config_file=config_file, + profile_name=profile, + ) @click.command(help='Deploy your lambda via S3.') @@ -138,6 +168,10 @@ def upload(use_requirements, local_package): default='config.yaml', help='Alternate config file.', ) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option( '--use-requirements', default=False, @@ -151,11 +185,13 @@ def upload(use_requirements, local_package): multiple=True, help='Install local package as well.', ) -def deploy_s3(use_requirements, local_package, config_file): +def deploy_s3(use_requirements, local_package, config_file, profile): aws_lambda.deploy_s3( - CURRENT_DIR, config_file=config_file, + CURRENT_DIR, use_requirements=use_requirements, local_package=local_package, + config_file=config_file, + profile_name=profile, ) @@ -165,14 +201,21 @@ def deploy_s3(use_requirements, local_package, config_file): default='config.yaml', help='Alternate config file.', ) +@click.option( + '--profile', + help='AWS profile to use.', +) @click.option( '--keep-last', type=int, prompt='Please enter the number of recent versions to keep', ) -def cleanup(keep_last, config_file): +def cleanup(keep_last, config_file, profile): aws_lambda.cleanup_old_versions( - CURRENT_DIR, keep_last, config_file=config_file, + CURRENT_DIR, + keep_last, + config_file=config_file, + profile_name=profile, )