#!/usr/bin/env python3 | |
# Purpose: Create EMR bootstrap script bucket and deploy the cfn stack | |
# Author: Gary A. Stafford (December 2020) | |
# Reference: https://gist.github.com/svrist/73e2d6175104f7ab4d201280acba049c | |
# Usage Example: python3 ./create_cfn_stack.py \ | |
# --ec2-key-name emr-demo-123456789012-us-east-1 \ | |
# --ec2-subnet-id subnet-06aa61f790a932b32 \ | |
# --environment dev | |
import argparse | |
import json | |
import logging | |
import os | |
import boto3 | |
from botocore.exceptions import ClientError | |
sts_client = boto3.client('sts') | |
cfn_client = boto3.client('cloudformation') | |
region = boto3.DEFAULT_SESSION.region_name | |
s3_client = boto3.client('s3', region_name=region) | |
logging.basicConfig(format='[%(asctime)s] %(levelname)s - %(message)s', level=logging.INFO) | |
def main(): | |
args = parse_args() | |
# create bootstrap bucket | |
account_id = sts_client.get_caller_identity()['Account'] | |
bootstrap_bucket = f'superset-emr-demo-bootstrap-{account_id}-{region}' | |
create_bucket(bootstrap_bucket) | |
# upload bootstrap script | |
dir_path = os.path.dirname(os.path.realpath(__file__)) | |
upload_file(f'{dir_path}/bootstrap_emr/bootstrap_actions.sh', bootstrap_bucket, 'bootstrap_actions.sh') | |
# set variables | |
stack_name = f'emr-superset-demo-{args.environment}' | |
cfn_template_path = f'{dir_path}/cloudformation/superset-emr-demo.yml' | |
cfn_params_path = f'{dir_path}/cloudformation/superset-emr-demo-params-{args.environment}.json' | |
ec2_key_name = args.ec2_key_name | |
# append new parameters | |
cfn_params = _parse_parameters(cfn_params_path) | |
cfn_params.append({'ParameterKey': 'Ec2KeyName', 'ParameterValue': ec2_key_name}) | |
cfn_params.append({'ParameterKey': 'Ec2SubnetId', 'ParameterValue': args.ec2_subnet_id}) | |
cfn_params.append({'ParameterKey': 'BootstrapBucket', 'ParameterValue': bootstrap_bucket}) | |
logging.info(json.dumps(cfn_params, indent=4)) | |
# create the cfn stack | |
create_stack(stack_name, cfn_template_path, cfn_params) | |
def create_bucket(bootstrap_bucket): | |
"""Create an S3 bucket in a specified region | |
:param bootstrap_bucket: Bucket to create | |
:return: True if bucket created, else False | |
""" | |
try: | |
s3_client.create_bucket(Bucket=bootstrap_bucket) | |
logging.info(f'New bucket name: {bootstrap_bucket}') | |
except ClientError as e: | |
logging.error(e) | |
return False | |
return True | |
def upload_file(file_name, bootstrap_bucket, object_name): | |
"""Upload a file to an S3 bucket | |
:param file_name: File to upload | |
:param bootstrap_bucket: Bucket to upload to | |
:param object_name: S3 object name | |
:return: True if file was uploaded, else False | |
""" | |
# Upload the file | |
try: | |
response = s3_client.upload_file(file_name, bootstrap_bucket, object_name) | |
logging.info(f'File {file_name} uploaded to bucket {bootstrap_bucket} as object {object_name}') | |
except ClientError as e: | |
logging.error(e) | |
return False | |
return True | |
def create_stack(stack_name, cfn_template, cfn_params): | |
"""Create EMR Cluster CloudFormation stack""" | |
template_data = _parse_template(cfn_template) | |
create_stack_params = { | |
'StackName': stack_name, | |
'TemplateBody': template_data, | |
'Parameters': cfn_params, | |
'TimeoutInMinutes': 60, | |
'Capabilities': [ | |
'CAPABILITY_NAMED_IAM', | |
], | |
'Tags': [ | |
{ | |
'Key': 'Project', | |
'Value': 'Superset EMR Demo' | |
}, | |
] | |
} | |
try: | |
response = cfn_client.create_stack(**create_stack_params) | |
logging.info(f'Response: {response}') | |
except ClientError as e: | |
logging.error(e) | |
return False | |
return True | |
def _parse_template(template): | |
with open(template) as template_file_obj: | |
template_data = template_file_obj.read() | |
cfn_client.validate_template(TemplateBody=template_data) | |
return template_data | |
def _parse_parameters(parameters): | |
with open(parameters) as parameter_file_obj: | |
parameter_data = json.load(parameter_file_obj) | |
return parameter_data | |
def parse_args(): | |
"""Parse argument values from command-line""" | |
parser = argparse.ArgumentParser(description='Arguments required for script.') | |
parser.add_argument('-e', '--environment', required=True, choices=['dev', 'test', 'prod'], help='Environment') | |
parser.add_argument('-k', '--ec2-key-name', required=True, help='Ec2KeyName: Name of EC2 Keypair') | |
parser.add_argument('-s', '--ec2-subnet-id', required=True, help='Ec2SubnetId: Name of EC2 Keypair') | |
args = parser.parse_args() | |
return args | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment