From aff275b68e68e7ebcc09dac6d85a6c6ab1e6e030 Mon Sep 17 00:00:00 2001
From: Terence Tuhinanshu <terence@tuhinanshu.com>
Date: Thu, 21 Dec 2023 10:14:30 -0500
Subject: [PATCH] Return task ARN from the handle command

When running ecsmanage in a programmatic setting, it is useful
to get the ARN of the task which it started, in addition to the
URL being printed to the console. This doesn't break existing
workflows, but adds the possibility of new ones.
---
 ecsmanage/management/commands/ecsmanage.py | 25 ++++++++++++----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/ecsmanage/management/commands/ecsmanage.py b/ecsmanage/management/commands/ecsmanage.py
index 2e1aa62..cc50b22 100644
--- a/ecsmanage/management/commands/ecsmanage.py
+++ b/ecsmanage/management/commands/ecsmanage.py
@@ -29,7 +29,7 @@ def add_arguments(self, parser):
     def handle(self, *args, **options):
         """
         Run the given command on the latest app CLI task definition and print
-        out a URL to view the status.
+        out a URL to view the status. Returns the task ARN of the started task.
         """
         self.env = options["env"]
         cmd = options["cmd"]
@@ -45,7 +45,15 @@ def handle(self, *args, **options):
         security_group_id = self.get_security_group(config["SECURITY_GROUP_TAGS"])
         subnet_id = self.get_subnet(config["SUBNET_TAGS"])
 
-        task_id = self.run_task(config, task_def_arn, security_group_id, subnet_id, cmd)
+        task_arn = self.run_task(config, task_def_arn, security_group_id, subnet_id, cmd)
+
+        # Task ARNs have at least two formats:
+        #
+        #  - Old: arn:aws:ecs:region:aws_account_id:task/task-id
+        #  - New: arn:aws:ecs:region:aws_account_id:task/cluster-name/task-id
+        #
+        # See: https://docs.aws.amazon.com/AmazonECS/latest/userguide/ecs-account-settings.html#ecs-resource-ids  # NOQA
+        task_id = task_arn.split("/")[-1]
 
         cluster_name = config["CLUSTER_NAME"]
 
@@ -58,6 +66,8 @@ def handle(self, *args, **options):
             self.style.SUCCESS(f"Task started! View here:\n{url}")
         )  # NOQA
 
+        return task_arn
+
     def parse_config(self):
         """
         Parse configuration settings for the app, checking to make sure that
@@ -165,7 +175,7 @@ def get_subnet(self, subnet_tags):
     def run_task(self, config, task_def_arn, security_group_id, subnet_id, cmd):
         """
         Run a task for a given task definition ARN using the given security
-        group and subnets, and return the task ID.
+        group and subnets, and return the task ARN of the started task.
         """
         task_def = self.ecs_client.describe_task_definition(
             taskDefinition=task_def_arn
@@ -200,11 +210,4 @@ def run_task(self, config, task_def_arn, security_group_id, subnet_id, cmd):
 
         task = self.parse_response(self.ecs_client.run_task(**kwargs), "tasks", 0)
 
-        # Task ARNs have at least two formats:
-        #
-        #  - Old: arn:aws:ecs:region:aws_account_id:task/task-id
-        #  - New: arn:aws:ecs:region:aws_account_id:task/cluster-name/task-id
-        #
-        # See: https://docs.aws.amazon.com/AmazonECS/latest/userguide/ecs-account-settings.html#ecs-resource-ids  # NOQA
-        task_id = task["taskArn"].split("/")[-1]
-        return task_id
+        return task["taskArn"]