add logic to wait for job completion

tools/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/amplify.go

@@ -15,8 +15,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 +191,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/main.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() {
@@ -96,6 +104,17 @@ func main() {
 		}
 	}
 
+	for i := 0; *wait && !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)
@@ -104,7 +123,7 @@ func main() {
 	logger.Info("Successfully posted PR comment")
 
 	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)
 	}
 }