/
lambda.py
178 lines (156 loc) · 7.23 KB
/
lambda.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
import boto3
import collections
import json
import logging
import os
DEFAULT_LOG_LEVEL = logging.DEBUG
REGION = os.environ['REGION']
ACCOUNT_ID = os.environ['ACCOUNT_ID']
sts_client = boto3.client('sts')
LOG_LEVELS = collections.defaultdict(
lambda: DEFAULT_LOG_LEVEL,
{
"critical": logging.CRITICAL,
"error": logging.ERROR,
"warning": logging.WARNING,
"info": logging.INFO,
"debug": logging.DEBUG,
},
)
# Lambda initializes a root logger that needs to be removed in order to set a
# different logging config
root = logging.getLogger()
if root.handlers:
for handler in root.handlers:
root.removeHandler(handler)
logging.basicConfig(level=LOG_LEVELS[os.environ.get("LOG_LEVEL", "").lower()])
log = logging.getLogger(__name__)
# Retrieve Current Task definition
def retrieve_task_definition(events_client, cw_rule_name):
response = events_client.list_targets_by_rule(
Rule=cw_rule_name,
Limit=1
)
targets = response['Targets']
target = targets[0]
task_definition_arn = target['EcsParameters']['TaskDefinitionArn']
log.info('current task definition %s ', task_definition_arn)
return task_definition_arn, target
# Retrieve Current Image url from current task definition
def retrieve_current_image(ecs_client, task_definition_arn, service):
response = ecs_client.describe_task_definition(
taskDefinition=task_definition_arn
)
task_definition = response['taskDefinition']
container_definitions = task_definition['containerDefinitions']
log.info(container_definitions)
# required if task definition has multiple containers like sidecars
container_definition = find_container(container_definitions, service)
if container_definition is None:
message = 'Couldnt find any image containing ' + service + ' in current taskDefinition ' \
+ task_definition_arn + '. Aborting '
raise Exception(message)
log.info('Current image definition: %s', container_definition['image'])
return container_definition['image'], task_definition
# Create new task definition using the old task definition and deployment image
def register_new_task_definition(ecs_client, task_definition, current_image, deployment_image):
container_definitions = task_definition['containerDefinitions']
# required if task definition has multiple containers like sidecars
container_definition = find_container(container_definitions, current_image)
container_definition['image'] = deployment_image
log.info('New Task def %s', task_definition)
response = ecs_client.register_task_definition(
family=task_definition['family'],
taskRoleArn=task_definition['taskRoleArn'],
executionRoleArn=task_definition['executionRoleArn'],
networkMode=task_definition['networkMode'],
containerDefinitions=container_definitions,
volumes=task_definition['volumes'],
placementConstraints=task_definition['placementConstraints'],
requiresCompatibilities=task_definition['requiresCompatibilities'],
cpu=task_definition['cpu'],
memory=task_definition['memory']
)
new_task_def = response['taskDefinition']['taskDefinitionArn']
log.info('New task definition: %s', new_task_def)
return new_task_def
# update cw rules target with the new task definition
def update_cw_rule_target(events_client, cw_rule_name, target, new_deployment_task_arn):
target['EcsParameters']['TaskDefinitionArn'] = new_deployment_task_arn
response = events_client.put_targets(
Rule=cw_rule_name,
Targets=[
target
]
)
log.info('Update Response: %s', response)
http_status = response['ResponseMetadata']['HTTPStatusCode']
if http_status != 200:
message = 'Unable to update target for cw rule: ' + cw_rule_name
raise Exception(message)
return
def find_container(containers, service):
for x in containers:
if service in x['image']:
return x
# sometimes tasks have different name but are using a same image, so service name doesn't match the image
# in this scenario return the first container definition
if len(containers) > 0:
return containers[0]
return None
def handler(event, context):
log.info("Received event: %s", json.dumps(event))
service = event['service']
deployment_image = event['image']
assume_role = event['assumeRole']
cw_rule_name = event['cwRuleName']
# Assume Role
log.info('Assuming role : %s', assume_role)
assumed_role = sts_client.assume_role(RoleArn=assume_role,
RoleSessionName="AssumeRoleSession1",
DurationSeconds=1800)
log.info("Successfully assumed role %s", assume_role)
events_client = boto3.client(
'events',
region_name=REGION,
aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
aws_session_token=assumed_role['Credentials']['SessionToken']
)
# Describe Service to retrieve Task Definition
ecs_client = boto3.client(
'ecs',
region_name=REGION,
aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
aws_session_token=assumed_role['Credentials']['SessionToken']
)
# Retrieve Task Definition
log.info('Finding target of rule of scheduled task cloudwatch rule of %s', cw_rule_name)
current_task_definition_arn, cw_target = retrieve_task_definition(events_client, cw_rule_name)
# Get Current Image
log.info('Retrieving current image name for service in taskDefinition: %s', current_task_definition_arn)
current_image, current_task_definition = retrieve_current_image(ecs_client, current_task_definition_arn,
service)
# compare current image with existing image to check if new deployment is needed
log.info('Comparing %s with %s', current_image, deployment_image)
if current_image == deployment_image:
log.info('image matches no new deployment needed')
event['previousImage'] = current_image
event['previousTaskDefArn'] = current_task_definition_arn
event['deployedTaskDefArn'] = current_task_definition_arn
event['deploymentNeeded'] = False
else:
log.info('image dos not match new deployment needed')
log.info('Creating new Task definition using %s', deployment_image)
new_deployment_task_arn = register_new_task_definition(ecs_client, current_task_definition, current_image,
deployment_image)
log.info('Updating Cloudwatch event target of %s with new task definition %s', cw_rule_name,
new_deployment_task_arn)
update_cw_rule_target(events_client, cw_rule_name, cw_target, new_deployment_task_arn)
event['previousImage'] = current_image
event['previousTaskDefArn'] = current_task_definition_arn
event['deployedTaskDefArn'] = new_deployment_task_arn
event['deploymentNeeded'] = True
log.info("Output: %s", json.dumps(event))
return event