def add_secrets(containers: List[dict], secrets: dict) -> List[dict]:
kctx = kitipy.get_current_context()
by_name = {c["name"]: c for c in containers}
for container, container_secrets in secrets.items():
by_name[container].update({"secrets": container_secrets})
return list(by_name.values())
def wait_for(tester: TesterCallable,
max_checks: int,
interval: int = 1,
label: Optional[str] = None):
"""This helper function will run a tester callback for max_checks times at
most, with an interval between each check. If the callback didn't return
true or a successful subprocess.CompletedProcess after max_checks retries,
an exception is thrown.
This helps implementing some higher-level wait functions, for instance to
ensure containers are all running or to ensure a DB is initialized.
Args:
tester (TesterCallable):
The callback to regularly run.
max_checks (int):
Number of times the tester functions should be called at most.
After that, an exception is thrown if no check were successful.
interval (int):
Interval in seconds between two retry.
label (Optional[str]):
Label to display on the CLI every time the tester function is
called. This is prefixed by "[X/max_checks]". It's also used for
the exception message when wait_for fails. It's recommended to
write it in the form of: "Wait for <something>".
Raises:
click.ClickException: When max_checks is reached and no checks were successful.
"""
kctx = kitipy.get_current_context()
label = label if label is not None else 'Waiting...'
for i in range(1, max_checks, interval):
kctx.echo(message="[%d/%d] %s" % (i, max_checks, label))
result = None
succeeded = False
try:
result = tester(kctx)
except subprocess.CalledProcessError as e:
succeedded = False
if isinstance(result, bool):
succeeded = result
if isinstance(result, subprocess.CompletedProcess):
succeeded = result.returncode == 0
if succeeded:
return
time.sleep(interval)
kctx.fail("Failed to %s" % (label.lower()))
def run_oneoff_task(client: mypy_boto3_ecs.ECSClient, cluster_name: str,
task_name: str, task_def: dict, container: str,
command: List[str], run_args: dict) -> str:
"""Run a specific command in a oneoff ECS task.
Args:
client (mypy_boto3_ecs.ECSClient):
An ECS API client.
cluster_name (str):
The name of the cluster where the task should run.
task_name (str):
The name of the task to create.
task_def (dict):
The task definition to register and deploy. See https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_definition_parameters.html.
container (str):
The name of the container where the command should run.
command (List[str]):
The shell command to run in the container.
run_args (dict):
The list of arguments to pass to run_task(). See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.run_task.
Returns:
str: The ARN of the task.
"""
# @TODO: use a proper logger
kctx = kitipy.get_current_context()
task_def_id = register_task_definition(client, task_def)
run_args["cluster"] = cluster_name
run_args["group"] = task_name
run_args["taskDefinition"] = task_def_id
run_args["count"] = 1
run_args["overrides"] = {
"containerOverrides": [{
"name": container,
"command": command
}]
}
resp = client.run_task(**run_args)
task_arn = resp["tasks"][0]["taskArn"]
kctx.info("A new oneoff task {0} has been scheduled.".format(task_arn))
return task_arn
def register_task_definition(client: mypy_boto3_ecs.ECSClient,
task_def: dict) -> str:
"""Register a task definition and returns its id.
Args:
client (mypy_boto3_ecs.ECSClient):
An ECS API client.
task_def (dict):
A task definition as expected by ECS API. See https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_definition_parameters.html.
Returns:
str: The definition ID in format "family:revision".
"""
resp = client.register_task_definition(**task_def)
task_def_id = "{0}:{1}".format(resp["taskDefinition"]["family"],
resp["taskDefinition"]["revision"])
# @TODO: use a proper logger
kctx = kitipy.get_current_context()
kctx.info(("A new task definition {task_def_id} " +
"has been registered").format(task_def_id=task_def_id))
return task_def_id
def watch_deployment(
client: mypy_boto3_ecs.ECSClient,
cluster_name: str,
service_name: str,
deployment_id: str,
max_attempts: int = 120,
) -> Generator[mypy_boto3_ecs.type_defs.ServiceEventTypeDef, None, None]:
"""Wait until a service deployment is complete and stream ECS events.
This function polls the ECS API every 5s until the given deployment has
completed. A deployment is completed once it has PRIMARY status and its
number of desired replicas matches the running count.
Args:
client (mypy_boto3_ecs.ECSClient):
An ECS API client.
cluster_name (str):
The name of the cluster where the service run.
service_name (str):
The name of the service to look for.
deployment_id (str):
The ID of the deployment to watch.
max_attempts (number):
The maximum number of attempts to be made. Default: 120 (~10 minutes).
Raises:
ServiceNotFoundError: When no matching service was found.
RuntimeError: When more than 1 service have been returned by ECS API.
DeploymentNotFoundError: When no deployment with the given ID is found.
RuntimeError: When max_attempts is reached.
"""
kctx = kitipy.get_current_context()
status = None
last_date = None
attempts = 0.
while attempts < max_attempts:
deployment = find_service_deployment(
client, cluster_name, service_name,
lambda d: d["id"] == deployment_id)
if deployment is None:
raise DeploymentNotFoundError(
"Deployment {0} not found.".format(deployment_id))
if last_date is None:
last_date = deployment["createdAt"]
status = deployment["status"]
events = list_service_events(client, cluster_name, service_name)
new_events = list(e for e in events if e["createdAt"] > last_date)
if len(new_events) > 0:
last_date = new_events[0]["createdAt"]
for event in reversed(new_events):
yield event
running_count = deployment["runningCount"]
desired_count = deployment["desiredCount"]
if status == "PRIMARY" and running_count == desired_count:
return
time.sleep(5)
attempts += 1
raise RuntimeError(
"watch_deployment timed out before the deployment was completed. It is probably broken."
)
def upsert_service(client: mypy_boto3_ecs.ECSClient, cluster_name: str,
service_name: str, task_def: dict,
service_def: dict) -> str:
"""Upsert an ECS service with its task definition.
The desiredCount of the current service deployment is automatically reused.
Args:
client (mypy_boto3_ecs.ECSClient):
An ECS API client.
cluster_name (str):
The name of the cluster where the service should be looked for.
service_name (str):
The name of the service to look for.
task_def (dict):
The task definition to register and deploy. See https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_definition_parameters.html.
service_def (dict):
The definition of the service to upsert. See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.create_service.
Returns:
string: The ID of the service deployment.
Raises:
ServiceDefinitionChangedError:
Both loadBalancers and serviceRegistries parameters from the
service definitions can't be changed after creation. If you need to
update these parameters, you should change the service name.
"""
# @TODO: use a proper logger
kctx = kitipy.get_current_context()
task_def_id = register_task_definition(client, task_def)
service_def["cluster"] = cluster_name
service_def["serviceName"] = service_name
service_def["taskDefinition"] = task_def_id
if find_service_arn(client, cluster_name, service_name) is None:
kctx.info(("Creating service {service} " +
"in {cluster} cluster.").format(service=service_name,
cluster=cluster_name))
resp = client.create_service(**service_def)
return resp["service"]["deployments"][0]["id"]
existing = describe_service(client, cluster_name, service_name)
if existing["loadBalancers"] != service_def.get("loadBalancers", []):
raise ServiceDefinitionChangedError(
"The parameter loadBalancers has changed.")
if existing["serviceRegistries"] != service_def.get(
"serviceRegistries", []):
# @TODO: add previous/current values to the exception
raise ServiceDefinitionChangedError(
"The parameter serviceRegistries has changed.")
# Remvoe all the params that are supported by create_service but not by
# update_service.
service_def["service"] = service_def["serviceName"]
service_def = {
k: v
for k, v in service_def.items() if k not in create_update_diff
}
kctx.info(("Updating service {service} " + "in {cluster} cluster.").format(
service=service_name, cluster=cluster_name))
service_def["desiredCount"] = existing["desiredCount"]
resp = client.update_service(**service_def)
return resp["service"]["deployments"][0]["id"]