deploy-ecs.main.kts

@file:DependsOn("aws.sdk.kotlin:ssm-jvm:0.28.2-beta")
@file:DependsOn("aws.sdk.kotlin:ecs-jvm:0.28.2-beta")

import aws.sdk.kotlin.services.ecs.*
import aws.sdk.kotlin.services.ecs.model.*
import aws.sdk.kotlin.services.ssm.*
import kotlinx.coroutines.*
import kotlin.system.exitProcess
import kotlin.time.Duration
import kotlin.time.Duration.Companion.minutes

if (args.size != 1) {
    println("Usage: kotlin deploy.main.kts <docker tag>")
    exitProcess(1)
}
val dockerImageTag = args[0]
val parameterName = "/prod/backend/docker-image-tag"

val ecs = Ecs()
val paramStore = ParamStore()

if (paramStore.read(parameter = parameterName) == dockerImageTag) {
    println("Docker image with tag $dockerImageTag is already deployed")
    exitProcess(0)
}

val prodCluster = ecs.cluster(name = "prod")
val backendService = prodCluster.service(name = "backend")
val currentTaskDefinition = backendService.taskDefinition()

val newTaskDefinition = currentTaskDefinition.clone {
    containerDefinitions = containerDefinitions?.replaceDockerImageTag(container = "backend", tag = dockerImageTag)
}.createNewRevision()

val deployment = backendService.setTaskDefinition(newTaskDefinition)
deployment.waitCompleted(timeout = 10.minutes)

paramStore.write(parameter = parameterName, value = dockerImageTag)

ecs.close()
paramStore.close()

////////////////////////////////////////////////////////

class Ecs {
    private val client: EcsClient

    init {
        client = runBlocking { EcsClient.fromEnvironment() }
        println("Region: ${client.config.region}")
    }

    fun cluster(name: String) = Cluster(client, name)
    fun close() = client.close()
}

class Cluster(
    private val client: EcsClient,
    name: String,
) {
    private val cluster: aws.sdk.kotlin.services.ecs.model.Cluster

    init {
        cluster = runBlocking {
            val arn = client.listClusters().clusterArns
                ?.single { it.endsWith(":cluster/$name") }
                ?: error("Cluster '$name' not found")
            client.describeClusters { clusters = listOf(arn) }.clusters
                ?.single { it.clusterName == name }
                ?: error("Cluster '$name' not found")
        }
        println("Cluster ARN: ${cluster.clusterArn}")
    }

    fun service(name: String) = Service(client, cluster, name)
}

class Service(
    private val client: EcsClient,
    private val cluster: aws.sdk.kotlin.services.ecs.model.Cluster,
    name: String,
) {
    private val service: aws.sdk.kotlin.services.ecs.model.Service

    init {
        service = runBlocking {
            val svc = client.describeServices {
                cluster = this@Service.cluster.clusterArn
                services = listOf(name)
            }.services ?: error("Service list is empty")
            svc.single()
        }
        println("Service ARN: ${service.serviceArn}")
    }

    fun taskDefinition() = TaskDefinition(client, service.taskDefinition)

    fun setTaskDefinition(task: TaskDefinition): Deployment {
        val deployments = runBlocking {
            client.updateService {
                cluster = this@Service.cluster.clusterArn
                this.service = this@Service.service.serviceName
                taskDefinition = task.arn
            }

            client.describeServices {
                cluster = this@Service.cluster.clusterArn
                services = listOf(service.serviceName ?: error("Service name is empty"))
            }.services?.single()?.deployments ?: error("Deployment list is empty")
        }
        return Deployment(client, cluster, service, deployments.single { it.status == "PRIMARY" })
    }
}


class TaskDefinition {
    private val client: EcsClient
    private val taskDefinition: aws.sdk.kotlin.services.ecs.model.TaskDefinition
    private val tags: List<Tag>

    constructor(client: EcsClient, arn: String?) {
        this.client = client
        val task = runBlocking {
            client.describeTaskDefinition {
                taskDefinition = arn ?: error("Task definition ARN is empty")
                include = listOf(TaskDefinitionField.Tags)
            }
        }
        taskDefinition = task.taskDefinition ?: error("Task definition is empty")
        tags = task.tags ?: emptyList()
        println("Current task definition ARN: ${taskDefinition.taskDefinitionArn}")
    }

    constructor(client: EcsClient, taskDefinition: aws.sdk.kotlin.services.ecs.model.TaskDefinition, tags: List<Tag>) {
        this.client = client
        this.taskDefinition = taskDefinition
        this.tags = tags
    }

    val arn: String
        get() = taskDefinition.taskDefinitionArn ?: error("Task definition ARN is empty")

    fun clone(block: aws.sdk.kotlin.services.ecs.model.TaskDefinition.Builder.() -> Unit): TaskDefinition {
        return TaskDefinition(client, taskDefinition.copy(block), tags)
    }

    fun createNewRevision(): TaskDefinition {
        val task = runBlocking {
            client.registerTaskDefinition {
                family = taskDefinition.family
                requiresCompatibilities = taskDefinition.requiresCompatibilities
                taskRoleArn = taskDefinition.taskRoleArn
                executionRoleArn = taskDefinition.executionRoleArn
                networkMode = taskDefinition.networkMode
                runtimePlatform = taskDefinition.runtimePlatform
                cpu = taskDefinition.cpu
                memory = taskDefinition.memory
                containerDefinitions = taskDefinition.containerDefinitions

                placementConstraints = taskDefinition.placementConstraints
                volumes = taskDefinition.volumes
                ephemeralStorage = taskDefinition.ephemeralStorage
                ipcMode = taskDefinition.ipcMode
                pidMode = taskDefinition.pidMode

                this.tags = this@TaskDefinition.tags
            }.taskDefinition ?: error("Task definition is empty")
        }

        println("Registered task definition ARN: ${task.taskDefinitionArn}")
        return TaskDefinition(client, task, tags)
    }
}

class Deployment(
    private val client: EcsClient,
    private val cluster: aws.sdk.kotlin.services.ecs.model.Cluster,
    private val service: aws.sdk.kotlin.services.ecs.model.Service,
    private val deployment: aws.sdk.kotlin.services.ecs.model.Deployment,
) {
    init {
        println("Deployment ID: ${deployment.id}")
    }

    fun waitCompleted(timeout: Duration) {
        runBlocking {
            withTimeout(timeout.inWholeMilliseconds) {
                while (true) {
                    val deployments = client.describeServices {
                        cluster = this@Deployment.cluster.clusterArn
                        services = listOf(service.serviceName ?: error("Service name is empty"))
                    }.services?.single()?.deployments ?: error("Deployment list is empty")
                    val deployment = deployments.single { it.id == deployment.id }
                    if (deployment.rolloutState == DeploymentRolloutState.Completed) {
                        break
                    }
                    println("Deployment state: ${deployment.rolloutState}")
                    delay(5_000)
                }
            }
        }
        println("Done!")
    }
}

fun List<ContainerDefinition>.replaceDockerImageTag(container: String, tag: String): List<ContainerDefinition> {
    println("Replacing Docker tag for '$container' to '$tag'")
    return map {
        if (it.name == container) {
            println("Current image: ${it.image}")
            it.copy {
                val parts = it.image?.split(":")
                val image = if (parts?.size == 2) {
                    parts[0]
                } else {
                    error("Image does not contain tag")
                }
                val s = "$image:$tag"
                println("New image: $s")
                this.image = s
            }
        } else {
            it
        }
    }
}

class ParamStore {
    private val client: SsmClient

    init {
        client = runBlocking { SsmClient.fromEnvironment() }
    }

    fun read(parameter: String): String {
        val tag = runBlocking {
            client.getParameter { this.name = parameter }.parameter?.value ?: error("Parameter value is empty")
        }
        println("Read from Parameter store: $parameter=$tag")
        return tag
    }

    fun write(parameter: String, value: String) {
        val param = runBlocking {
            client.putParameter {
                this.name = parameter
                this.value = value
                overwrite = true
            }
        }
        println("Parameter store updated: $parameter=$value (version ${param.version})")
    }

    fun close() = client.close()
}