Enable run_task command to wait until the executions ends and return …

…the exit code CircleCI-Public#230
gonrial · Dec 10, 2024 · 5c381e7 · 5c381e7
1 parent 1c53fd3
commit 5c381e7
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 6 deletions.
diff --git a/.circleci/test-deploy.yml b/.circleci/test-deploy.yml
@@ -439,7 +439,7 @@ workflows:
           aws-resource-name-prefix: ${AWS_RESOURCE_NAME_PREFIX_FARGATE}
           terraform-config-dir: "tests/terraform_setup/fargate"
           context: [CPE-OIDC]
-          role_arn: "arn:aws:iam::122211685980:role/CPE_ECS_OIDC_TEST"      
+          role_arn: "arn:aws:iam::122211685980:role/CPE_ECS_OIDC_TEST"
       - build-test-app:
           name: fargate_build-test-app
           docker-image-namespace: "${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_DEFAULT_REGION}.amazonaws.com"
@@ -574,11 +574,28 @@ workflows:
           run_task_output: "run_task_output.json"
           overrides: '{"containerOverrides":[{"name": "${INTERPOLATION_TEST}", "memory": 512}]}'
           context: [CPE-OIDC]
+      - aws-ecs/run_task:
+          name: ec2_run_task_and_wait-test
+          auth:
+            - aws-cli/setup:
+                role_arn: "arn:aws:iam::122211685980:role/CPE_ECS_OIDC_TEST"
+          filters: *filters
+          requires:
+            - ec2_set-up-run_task-test
+          cluster: "${AWS_RESOURCE_NAME_PREFIX_EC2}-cluster"
+          task_definition: "${AWS_RESOURCE_NAME_PREFIX_EC2}-sleep360"
+          launch_type: "EC2"
+          awsvpc: false
+          overrides: '{"containerOverrides":[{"name": "${INTERPOLATION_TEST}", "memory": 512}]}'
+          context: [CPE-OIDC]
+          wait_task_stopped: true
+          exit_code_from: sleep
       - tear-down-run_task-test:
           name: ec2_tear-down-run_task-test
           filters: *filters
           requires:
             - ec2_run_task-test
+            - ec2_run_task_and_wait-test
           family_name: ${AWS_RESOURCE_NAME_PREFIX_EC2}-sleep360
           context: [CPE-OIDC]
           role_arn: "arn:aws:iam::122211685980:role/CPE_ECS_OIDC_TEST"
@@ -846,7 +863,7 @@ commands:
             else
               PORT="<< parameters.port >>"
             fi
-            
+
             set -x
             TARGET_GROUP_ARN=$(aws ecs describe-services --cluster << parameters.cluster >> --services << parameters.service_name >> $PROFILE | jq -r '.services[0].loadBalancers[0].targetGroupArn')
             ELB_ARN=$(aws elbv2 describe-target-groups --target-group-arns $TARGET_GROUP_ARN $PROFILE | jq -r '.TargetGroups[0].LoadBalancerArns[0]')

diff --git a/src/commands/run_task.yml b/src/commands/run_task.yml
@@ -128,6 +128,14 @@ parameters:
     description: AWS region to use. Defaults to AWS_DEFAULT_REGION environment variable.
     type: string
     default: $AWS_DEFAULT_REGION
+  wait_task_stopped:
+    description: Wait until the task execution ends. Doesn't work with `run_task_output`.
+    type: boolean
+    default: false
+  exit_code_from:
+    description: Name of the container which exit code will be returned if wait_task_stopped is true. Defaults to the first container.
+    type: string
+    default: ''
 steps:
   - run:
       name: Run Task
@@ -155,3 +163,5 @@ steps:
         ORB_STR_CD_CAPACITY_PROVIDER_STRATEGY: <<parameters.capacity_provider_strategy>>
         ORB_STR_PROFILE_NAME: <<parameters.profile_name>>
         ORB_STR_RUN_TASK_OUTPUT: <<parameters.run_task_output>>
+        ORB_BOOL_WAIT_TASK_STOPPED: <<parameters.wait_task_stopped>>
+        ORB_STR_EXIT_CODE_FROM: <<parameters.exit_code_from>>
diff --git a/src/examples/run_task_and_wait.yml b/src/examples/run_task_and_wait.yml
@@ -0,0 +1,27 @@
+description: Start the execution of an ECS task and wait until the task is completed.
+usage:
+  version: 2.1
+  orbs:
+    aws-ecs: circleci/[email protected]
+    aws-cli: circleci/[email protected]
+  jobs:
+    run_task:
+      docker:
+        - image: cimg/python:3.10
+      steps:
+        - aws-cli/setup:
+            profile_name: "OIDC-USER"
+            role_arn: "arn:aws:iam::123456789012:role/VALID_OIDC_ECS_ROLE"
+        - aws-ecs/run_task:
+            cluster: $CLUSTER_NAME
+            launch_type: ""
+            region: us-east-1
+            task_definition: $My_Task_Def
+            subnet_ids: '$SUBNET_ONE, $SUBNET_TWO'
+            security_group_ids: $SECURITY_GROUP_IDS_FETCHED\
+            wait_task_stopped: true
+            exit_code_from: app
+  workflows:
+    run_task:
+      jobs:
+        - run_task
diff --git a/src/scripts/run_task.sh b/src/scripts/run_task.sh
@@ -14,6 +14,7 @@ ORB_STR_RUN_TASK_OUTPUT="$(circleci env subst "$ORB_STR_RUN_TASK_OUTPUT")"
 ORB_STR_PROFILE_NAME="$(circleci env subst "$ORB_STR_PROFILE_NAME")"
 ORB_STR_ASSIGN_PUB_IP="$(circleci env subst "$ORB_STR_ASSIGN_PUB_IP")"
 ORB_AWS_REGION="$(circleci env subst "$ORB_AWS_REGION")"
+ORB_STR_EXIT_CODE_FROM="$(circleci env subst "$ORB_STR_EXIT_CODE_FROM")"
 
 if [[ "$ORB_STR_OVERRIDES" == *"\${"* ]]; then
     ORB_STR_OVERRIDES="$(echo "${ORB_STR_OVERRIDES}" | circleci env subst)"
@@ -77,7 +78,6 @@ if [ -n "$ORB_STR_CD_CAPACITY_PROVIDER_STRATEGY" ]; then
     # shellcheck disable=SC2086
     set -- "$@" --capacity-provider-strategy $ORB_STR_CD_CAPACITY_PROVIDER_STRATEGY
 fi
-
 if [ -n "$ORB_VAL_LAUNCH_TYPE" ]; then
     if [ -n "$ORB_STR_CD_CAPACITY_PROVIDER_STRATEGY" ]; then
         echo "Error: "
@@ -88,7 +88,10 @@ if [ -n "$ORB_VAL_LAUNCH_TYPE" ]; then
         set -- "$@" --launch-type "$ORB_VAL_LAUNCH_TYPE"
     fi
 fi
-
+if [ "$ORB_BOOL_WAIT_TASK_STOPPED" == "1" ]; then
+    echo "Setting query to export taskArn"
+    set -- "$@" --query 'tasks[].taskArn' --output text
+fi
 
 echo "Setting --count"
 set -- "$@" --count "$ORB_INT_COUNT"
@@ -99,11 +102,52 @@ set -- "$@" --cluster "$ORB_STR_CLUSTER_NAME"
 
 
 if [ -n "${ORB_STR_RUN_TASK_OUTPUT}" ]; then
+    if [ "$ORB_BOOL_WAIT_TASK_STOPPED" == "1" ]; then
+        echo "Exporting the run_task_output parameter is not compatible with wait_task_stopped parameter."
+        exit 1
+    fi
+
     set -x
     aws ecs run-task --profile "${ORB_STR_PROFILE_NAME}" --region "${ORB_AWS_REGION}" "$@" | tee "${ORB_STR_RUN_TASK_OUTPUT}"
     set +x
 else
-    set -x    
-    aws ecs run-task --profile "${ORB_STR_PROFILE_NAME}" --region "${ORB_AWS_REGION}" "$@"
+    set -x
+    ORB_STR_TASK_ARN=$(aws ecs run-task --profile "${ORB_STR_PROFILE_NAME}" --region "${ORB_AWS_REGION}" "$@")
     set +x
 fi
+
+if [ "$ORB_BOOL_WAIT_TASK_STOPPED" == "1" ]; then
+    echo "Wait for ECS task $ORB_STR_TASK_ARN to stop..."
+
+    ORB_STR_WAIT_EXIT_CODE=$(aws ecs wait tasks-stopped \
+        --profile "${ORB_STR_PROFILE_NAME}" \
+        --region "${ORB_AWS_REGION}" \
+        --cluster "${ORB_STR_CLUSTER_NAME}" \
+        --tasks $ORB_STR_TASK_ARN
+    )
+
+    if [ "$ORB_STR_WAIT_EXIT_CODE" != "0" ]; then
+        echo "Stopped waiting for the task execution to end - please check the status of $ORB_STR_TASK_ARN on the AWS ECS console."
+    fi
+
+    # Get exit code
+    if [ -n "$ORB_STR_EXIT_CODE_FROM" ]
+        ORB_STR_TASK_EXIT_CODE=$(aws ecs describe-tasks \
+            --profile "${ORB_STR_PROFILE_NAME}" \
+            --region "${ORB_AWS_REGION}" \
+            --cluster "${ORB_STR_CLUSTER_NAME}" \
+            --tasks $TASK_ARN \
+            --query "tasks[0].containers[?name=='$ORB_STR_EXIT_CODE_FROM'].exitCode" \
+            --output text)
+    else
+        # Assume the first container
+        ORB_STR_TASK_EXIT_CODE=$(aws ecs describe-tasks \
+            --profile "${ORB_STR_PROFILE_NAME}" \
+            --region "${ORB_AWS_REGION}" \
+            --cluster "${ORB_STR_CLUSTER_NAME}" \
+            --tasks $TASK_ARN \
+            --query "tasks[0].containers[0].exitCode" \
+            --output text)
+    fi
+    exit $ORB_STR_TASK_EXIT_CODE
+fi