From 6871d376e46c057fc83015aaf1717fe4bda8666d Mon Sep 17 00:00:00 2001
From: Allan Denot <adenot@gmail.com>
Date: Wed, 21 Oct 2020 20:53:13 +1100
Subject: [PATCH] Fixes bug caused by task ARN coming with cluster name,
 causing script to not find logs

---
 src/run-task.sh       | 13 ++++++++-----
 src/tail-task-logs.py | 45 +++++++++++++++++++------------------------
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/run-task.sh b/src/run-task.sh
index d398f6c..a56cfa7 100755
--- a/src/run-task.sh
+++ b/src/run-task.sh
@@ -27,10 +27,10 @@ TASK_ARN=$(aws ecs register-task-definition \
   --output=text)
 
 echo "---> Executing  ECS Task"
-echo "       CLUSTER_NAME: ${CLUSTER_NAME}"
-echo "       APP_NAME: ${APP_NAME}"
-echo "       TASK_ARN: ${TASK_ARN}"
-
+echo "     CLUSTER_NAME: ${CLUSTER_NAME}"
+echo "     APP_NAME: ${APP_NAME}"
+echo "     TASK_DEFINITION_ARN: ${TASK_ARN}"
+echo -n "     STATUS: "
 TASK_ID=$(aws ecs run-task \
   --cluster $CLUSTER_NAME \
   --task-definition $TASK_ARN \
@@ -41,8 +41,11 @@ sleep 5
 
 while [ "$(aws ecs describe-tasks --tasks $TASK_ID --cluster $CLUSTER_NAME --query="tasks[0].lastStatus" --output=text)" == "PENDING" ]
 do
-  sleep 1
+  echo -n "."
+  sleep 5
 done
+echo -n "RUNNING"
+echo
 
 echo "---> Task ARN $TASK_ID"
 
diff --git a/src/tail-task-logs.py b/src/tail-task-logs.py
index ebffac9..c520c9c 100755
--- a/src/tail-task-logs.py
+++ b/src/tail-task-logs.py
@@ -9,20 +9,15 @@
 app_name=os.environ['APP_NAME']
 task_arn=sys.argv[1]
 
-task_id=task_arn.split(":task/",1)[1]  #get the task number id
+task_id=task_arn.split("/")[-1]  #get the task number id (without the cluster name)
 last_event = None
 log_group_name='/ecs/'+cluster_name+'/'+app_name
 log_stream_prefix = None
 
-print("     Waiting for logs...")
+print("======== TASK LOGS ========")
 
 while True:
   try:
-    response = aws_ecs.describe_tasks(
-      cluster=cluster_name,
-      tasks=[task_arn])
-    task_status = response['tasks'][0]['lastStatus']
-
     if log_stream_prefix is None:
       log_streams = logs.describe_log_streams(logGroupName=log_group_name, orderBy='LastEventTime', descending=True, limit=1)
 
@@ -33,7 +28,6 @@
           'logStreamName': log_stream_prefix+'/'+task_id,
           'startFromHead': True
         }
-
     else:
       log_stream_events = logs.get_log_events(**extra_args)
 
@@ -43,22 +37,24 @@
       if 'nextToken' not in extra_args or log_stream_events['nextForwardToken'] != extra_args['nextToken']:
         extra_args['nextToken'] = log_stream_events['nextForwardToken']
 
-    if task_status == "STOPPED":
-      print("======== TASK STOPPED ========")
-      print("Task ID:        %s" % task_id)
-      print("Task ARN:       %s" % task_arn)
-      print("Service Name:   %s" % app_name)
-      print("Cluster Name:   %s" % cluster_name)
-      if 'startedAt' in response['tasks'][0]:
-        print("Started at:     %s" % response['tasks'][0]['startedAt'])
-      print("Stopped at:     %s" % response['tasks'][0]['stoppedAt'])
-      print("Stopped Reason: %s" % response['tasks'][0]['stoppedReason'])
-      if 'stopCode' in response['tasks'][0]:
-        print("Stop Code:      %s" % response['tasks'][0]['stopCode'])
-      print("")
-      break
-
-    time.sleep(1)
+      response = aws_ecs.describe_tasks(
+        cluster=cluster_name,
+        tasks=[task_arn])
+
+      if response['tasks'][0]['lastStatus'] == "STOPPED":
+        print("======== TASK STOPPED ========")
+        print("Task ID:        %s" % task_id)
+        print("Task ARN:       %s" % task_arn)
+        print("Service Name:   %s" % app_name)
+        print("Cluster Name:   %s" % cluster_name)
+        if 'startedAt' in response['tasks'][0]:
+          print("Started at:     %s" % response['tasks'][0]['startedAt'])
+        print("Stopped at:     %s" % response['tasks'][0]['stoppedAt'])
+        print("Stopped Reason: %s" % response['tasks'][0]['stoppedReason'])
+        if 'stopCode' in response['tasks'][0]:
+          print("Stop Code:      %s" % response['tasks'][0]['stopCode'])
+        print("")
+        break
 
   except logs.exceptions.ResourceNotFoundException as e:
     time.sleep(5)
@@ -68,4 +64,3 @@
     print("Error: " + str(e))
     break
 
-