add logic to wait for job completion

Author: taraspos

.github/actions/amplify-preview/action.yaml

@@ -11,6 +11,9 @@ inputs:
   github_token:
     required: true
     description: "Github token with permissions to read/write comments in pull request"
+  wait:
+    default: "false"
+    description: "If Amplify deployment is pending/running state wait for it's completion"
 runs:
   using: composite
   steps:
@@ -30,6 +33,7 @@ runs:
         GIT_BRANCH_NAME: ${{ steps.extract_branch.outputs.branch }}
         CREATE_BRANCHES: ${{ inputs.create_branches }}
         GITHUB_TOKEN: ${{ inputs.github_token }}
+        WAIT: ${{ inputs.wait }} 
       shell: bash
       run: |
         pushd ./tools/amplify-preview/; go run ./; popd

tools/amplify-preview/action.yaml

@@ -1,3 +1,19 @@
+/*
+Copyright 2024 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package main
 
 import (
@@ -15,8 +31,13 @@ import (
 )
 
 var (
-	errBranchNotFound = errors.New("Branch not found")
-	errNoJobForBranch = errors.New("Current branch has no jobs")
+	errBranchNotFound           = errors.New("Branch not found")
+	errNoJobForBranch           = errors.New("Current branch has no jobs")
+	amplifyJobCompletedStatuses = map[types.JobStatus]struct{}{
+		types.JobStatusFailed:    {},
+		types.JobStatusCancelled: {},
+		types.JobStatusSucceed:   {},
+	}
 )
 
 const (
@@ -186,6 +207,11 @@ func appIDFromBranchARN(branchArn string) (string, error) {
 	return "", fmt.Errorf("Invalid branch ARN")
 }
 
+func isAmplifyJobCompleted(job *types.JobSummary) bool {
+	_, ok := amplifyJobCompletedStatuses[job.Status]
+	return ok
+}
+
 func (err aggregatedError) Error() error {
 	if len(err.perAppErr) == 0 {
 		return nil

tools/amplify-preview/amplify.go

@@ -1,3 +1,19 @@
+/*
+Copyright 2024 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package main
 
 import (

tools/amplify-preview/github.go

@@ -22,6 +22,7 @@ import (
 	"log/slog"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -38,6 +39,13 @@ var (
 	gitBranchName  = kingpin.Flag("git-branch-name", "Git branch name").Envar("GIT_BRANCH_NAME").Required().String()
 	createBranches = kingpin.Flag("crate-branches",
 		"Defines whether Amplify branches should be created if missing, or just lookup existing ones").Envar("CREATE_BRANCHES").Default("false").Bool()
+	wait = kingpin.Flag("wait",
+		"Wait for pending/running job to complete").Envar("wait").Default("false").Bool()
+)
+
+const (
+	jobWaitSleepTime    = 30 * time.Second
+	jobWaitTimeAttempts = 40
 )
 
 func main() {
@@ -100,11 +108,28 @@ func main() {
 		logger.Error("failed to post preview URL", "error", err)
 		os.Exit(1)
 	}
-
 	logger.Info("Successfully posted PR comment")
 
+	if *wait {
+		for i := 0; !isAmplifyJobCompleted(job) && i < jobWaitTimeAttempts; i++ {
+			job, err := amp.GetJob(ctx, branch, nil)
+			if err != nil {
+				logger.Error("failed to get amplify job", logKeyBranchName, *gitBranchName, "error", err)
+				os.Exit(1)
+			}
+
+			logger.Info("Job is not in a completed state yet. Sleeping...", logKeyBranchName, *gitBranchName, "job_status", job.Status, "job_id", job.JobId)
+			time.Sleep(jobWaitSleepTime)
+		}
+
+		if err := postPreviewURL(ctx, amplifyJobToMarkdown(job, branch)); err != nil {
+			logger.Error("failed to post preview URL", "error", err)
+			os.Exit(1)
+		}
+	}
+
 	if job.Status == types.JobStatusFailed {
-		logger.Error("amplify job is in failed state", "job_status", job.Status, "job_id", job.JobId)
+		logger.Error("amplify job is in failed state", logKeyBranchName, *gitBranchName, "job_status", job.Status, "job_id", job.JobId)
 		os.Exit(1)
 	}
 }