Fargate tasks¶

import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Literal

from cfn_typed_sumo.types import apigateway, dynamodb, ecs, elasticloadbalancingv2 as elbv2, events
from cfn_typed_sumo.types import iam, kinesisfirehose, lambda_, logs, route53, s3, sns, sqs

from cfn_typed import AnyStatement, Expr, Resources, Tag, allow, arn, assume_role_policy, get_att, join, policy_document, ref, sub
from cfn_typed.policy import ANY, arn_equals, aws, service, string_equals
from cfn_typed_sumo.constructs import topic_to_queue
from cfn_typed_sumo.patterns import queue_with_dlq

@dataclass(frozen=True)
class FargateTask:
    """What `fargate_task` returns: the resources, and what a Lambda needs to start tasks of it on demand."""

    name: str
    resources: Resources
    subnets: Expr[list[Expr[str]]]
    security_groups: list[Expr[str]]

    @property
    def cluster_arn(self) -> Expr[str]:
        return get_att(f"{self.name}Cluster", "Arn")

    @property
    def task_definition_arn(self) -> Expr[str]:
        return ref(f"{self.name}Task")

    @property
    def run_task_statements(self) -> list[AnyStatement]:
        """For the invoking Lambda's role: `RunTask` on this task definition, and `PassRole` for the task's two roles."""
        return [
            allow("ecs:RunTask", self.task_definition_arn, condition=arn_equals("ecs:cluster", self.cluster_arn)),
            allow("iam:PassRole", [get_att(f"{self.name}ExecutionRole", "Arn"), get_att(f"{self.name}TaskRole", "Arn")]),
        ]

    @property
    def environment(self) -> dict[str, Expr[str]]:
        """Variables for the invoking Lambda, so the template, not a string in the code, says where tasks run."""
        return {
            "CLUSTER_ARN": self.cluster_arn,
            "TASK_DEFINITION_ARN": self.task_definition_arn,
            "SUBNETS": join(",", self.subnets),
            "SECURITY_GROUPS": self.security_groups[0],
        }


def fargate_task(
    name: str,
    *,
    image: Expr[str],
    subnets: Expr[list[Expr[str]]],
    security_groups: list[Expr[str]],
    cpu: Literal["256", "512", "1024", "2048", "4096"] = "256",
    memory: Literal["512", "1024", "2048", "4096", "8192"] = "512",
    desired_count: int | None = 1,
    port: int | None = None,
    environment: Mapping[str, Expr[str]] | None = None,
    task_statements: Sequence[AnyStatement] = (),
    retention_days: int = 30,
    tags: list[Expr[Tag]] | None = None,
) -> FargateTask:
    """A cluster, a task definition and (with `desired_count`) a service running one container, with no load balancer.

    For a web service behind a load balancer see `fargate_service` in [CDK L2 constructs](cdk-l2.md). This one is
    for workers, and for tasks that a Lambda starts.

    Logical IDs: `{name}Cluster`, `{name}Logs`, `{name}ExecutionRole`, `{name}TaskRole`, `{name}Task` and
    `{name}Service`. The **execution role** is what ECS uses to pull the image and write logs, and the **task
    role** is what your code runs as, with `task_statements`. Mixing the two up is the usual mistake.
    Pass `desired_count=None` for a task that is only started on demand (`run_task_statements`).
    """
    logs_id, execution, task_role, task, svc, cluster = (
        f"{name}Logs", f"{name}ExecutionRole", f"{name}TaskRole", f"{name}Task", f"{name}Service", f"{name}Cluster",
    )
    container: ecs.TaskDefinitionContainerDefinition = {
        "Name": name,
        "Image": image,
        "Essential": True,
        "LogConfiguration": {"LogDriver": "awslogs", "Options": {
            "awslogs-group": ref(logs_id), "awslogs-region": ref("AWS::Region"), "awslogs-stream-prefix": name,
        }},
    }
    if port:
        container["PortMappings"] = [{"ContainerPort": port, "Protocol": "tcp"}]
    if environment:
        container["Environment"] = [{"Name": k, "Value": v} for k, v in environment.items()]
    task_props: ecs.TaskDefinitionProps = {
        "Family": name,
        "Cpu": cpu,
        "Memory": memory,
        "NetworkMode": "awsvpc",
        "RequiresCompatibilities": ["FARGATE"],
        "ExecutionRoleArn": get_att(execution, "Arn"),
        "TaskRoleArn": get_att(task_role, "Arn"),
        "ContainerDefinitions": [container],
    }
    role_props: iam.RoleProps = {"AssumeRolePolicyDocument": assume_role_policy(service("ecs-tasks.amazonaws.com"))}
    if task_statements:
        role_props["Policies"] = [iam.RolePolicyProperty(PolicyName=f"{name}-task", PolicyDocument=policy_document(task_statements))]
    cluster_props: ecs.ClusterProps = {"ClusterSettings": [{"Name": "containerInsights", "Value": "enabled"}]}
    _tags(task_props, tags)
    _tags(cluster_props, tags)
    resources: Resources = {
        cluster: ecs.cluster(**cluster_props),
        logs_id: logs.log_group(RetentionInDays=retention_days),
        execution: iam.role(
            AssumeRolePolicyDocument=assume_role_policy(service("ecs-tasks.amazonaws.com")),
            ManagedPolicyArns=[sub("arn:${AWS::Partition}:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy")],
        ),
        task_role: iam.role(**role_props),
        task: ecs.task_definition(**task_props),
    }
    if desired_count is not None:
        service_props: ecs.ServiceProps = {
            "Cluster": ref(cluster),
            "TaskDefinition": ref(task),
            "DesiredCount": desired_count,
            "LaunchType": "FARGATE",
            "NetworkConfiguration": {"AwsvpcConfiguration": {
                "Subnets": subnets, "SecurityGroups": security_groups, "AssignPublicIp": "DISABLED",
            }},
            "DeploymentConfiguration": {"DeploymentCircuitBreaker": {"Enable": True, "Rollback": True}},
        }
        _tags(service_props, tags)
        resources[svc] = ecs.service(**service_props)
    return FargateTask(name, resources, subnets, security_groups)

job = fargate_task(
    "Report", image=ref("ReportImage"), subnets=ref("PrivateSubnets"), security_groups=[ref("TaskGroup")],
    desired_count=None,                                    # started on demand, no service
    task_statements=[allow("s3:PutObject", arn.object("reports"))],
)
starter = lambda_role("starter", job.run_task_statements)  # and pass job.environment to the function