Address PR feedback

libs/github/comment.go

@@ -12,18 +12,21 @@ var (
 	ErrCommentNotFound = errors.New("comment not found")
 )
 
+// IssueIdentifier represents an issue or PR on GitHub
 type IssueIdentifier struct {
 	Owner  string
 	Repo   string
 	Number int
 }
 
-// every trait is treated as AND
+// CommentTraits defines optional traits to filter comments.
+// Every trait (if non-empty-string) is matched with an "AND" operator
 type CommentTraits struct {
-	BodyContains *string
-	UserLogin    *string
+	BodyContains string
+	UserLogin    string
 }
 
+// FindCommentByTraits searches for a comment in an PR or issue based on specified traits
 func (c *Client) FindCommentByTraits(ctx context.Context, issue IssueIdentifier, targetComment CommentTraits) (*github.IssueComment, error) {
 	comments, _, err := c.client.Issues.ListComments(ctx, issue.Owner, issue.Repo, issue.Number, nil)
 	if err != nil {
@@ -32,16 +35,12 @@ func (c *Client) FindCommentByTraits(ctx context.Context, issue IssueIdentifier,
 
 	for _, c := range comments {
 		matcher := true
-		if targetComment.UserLogin != nil {
-			matcher = matcher &&
-				c.User != nil && c.User.Login != nil &&
-				*c.User.Login == *targetComment.UserLogin
+		if targetComment.UserLogin != "" {
+			matcher = matcher && c.User.GetLogin() == targetComment.UserLogin
 		}
 
-		if targetComment.BodyContains != nil {
-			matcher = matcher &&
-				c.Body != nil &&
-				strings.Contains(*c.Body, *targetComment.BodyContains)
+		if targetComment.BodyContains != "" {
+			matcher = matcher && strings.Contains(c.GetBody(), targetComment.BodyContains)
 		}
 
 		if matcher {
@@ -52,6 +51,7 @@ func (c *Client) FindCommentByTraits(ctx context.Context, issue IssueIdentifier,
 	return nil, ErrCommentNotFound
 }
 
+// CreateComment creates a new comment on an issue or PR
 func (c *Client) CreateComment(ctx context.Context, issue IssueIdentifier, commentBody string) error {
 	_, _, err := c.client.Issues.CreateComment(ctx, issue.Owner, issue.Repo, issue.Number, &github.IssueComment{
 		Body: &commentBody,
@@ -60,6 +60,7 @@ func (c *Client) CreateComment(ctx context.Context, issue IssueIdentifier, comme
 	return err
 }
 
+// UpdateComment updates an existing comment on an issue or PR
 func (c *Client) UpdateComment(ctx context.Context, issue IssueIdentifier, commentId int64, commentBody string) error {
 	_, _, err := c.client.Issues.EditComment(ctx, issue.Owner, issue.Repo, commentId, &github.IssueComment{
 		Body: &commentBody,

tools/amplify-preview/amplify.go

@@ -51,16 +51,17 @@ const (
 )
 
 type AmplifyPreview struct {
-	appIDs []string
-	client *amplify.Client
+	appIDs     []string
+	branchName string
+	client     *amplify.Client
 }
 
 type aggregatedError struct {
 	perAppErr map[string]error
 	message   string
 }
 
-func (amp *AmplifyPreview) FindExistingBranch(ctx context.Context, branchName string) (*types.Branch, error) {
+func (amp *AmplifyPreview) FindExistingBranch(ctx context.Context) (*types.Branch, error) {
 	type resp struct {
 		appID string
 		data  *amplify.GetBranchOutput
@@ -75,7 +76,7 @@ func (amp *AmplifyPreview) FindExistingBranch(ctx context.Context, branchName st
 			defer wg.Done()
 			branch, err := amp.client.GetBranch(ctx, &amplify.GetBranchInput{
 				AppId:      aws.String(appID),
-				BranchName: aws.String(branchName),
+				BranchName: aws.String(amp.branchName),
 			})
 			resultCh <- resp{
 				appID: appID,
@@ -96,7 +97,7 @@ func (amp *AmplifyPreview) FindExistingBranch(ctx context.Context, branchName st
 	for resp := range resultCh {
 		var errNotFound *types.NotFoundException
 		if errors.As(resp.err, &errNotFound) {
-			logger.Debug("Branch not found", logKeyAppID, resp.appID, logKeyBranchName, branchName)
+			logger.Debug("Branch not found", logKeyAppID, resp.appID, logKeyBranchName, amp.branchName)
 			continue
 		} else if resp.err != nil {
 			failedResp.perAppErr[resp.appID] = resp.err
@@ -115,7 +116,7 @@ func (amp *AmplifyPreview) FindExistingBranch(ctx context.Context, branchName st
 	return nil, errBranchNotFound
 }
 
-func (amp *AmplifyPreview) CreateBranch(ctx context.Context, branchName string) (*types.Branch, error) {
+func (amp *AmplifyPreview) CreateBranch(ctx context.Context) (*types.Branch, error) {
 	failedResp := aggregatedError{
 		perAppErr: map[string]error{},
 		message:   "failed to create branch",
@@ -124,7 +125,7 @@ func (amp *AmplifyPreview) CreateBranch(ctx context.Context, branchName string)
 	for _, appID := range amp.appIDs {
 		resp, err := amp.client.CreateBranch(ctx, &amplify.CreateBranchInput{
 			AppId:           aws.String(appID),
-			BranchName:      aws.String(branchName),
+			BranchName:      aws.String(amp.branchName),
 			Description:     aws.String("Branch generated for PR TODO"),
 			Stage:           types.StagePullRequest,
 			EnableAutoBuild: aws.Bool(true),
@@ -169,38 +170,87 @@ func (amp *AmplifyPreview) StartJob(ctx context.Context, branch *types.Branch) (
 
 }
 
-func (amp *AmplifyPreview) GetLatestAndActiveJobs(ctx context.Context, branch *types.Branch) (latestJob, activeJob *types.JobSummary, err error) {
+func (amp *AmplifyPreview) findJobsByID(ctx context.Context, branch *types.Branch, includeLatest bool, ids ...string) (jobSummaries []types.JobSummary, err error) {
 	appID, err := appIDFromBranchARN(*branch.BranchArn)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	resp, err := amp.client.ListJobs(ctx, &amplify.ListJobsInput{
 		AppId:      aws.String(appID),
-		BranchName: branch.BranchName,
+		BranchName: aws.String(amp.branchName),
 		MaxResults: 50,
 	})
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 	if len(resp.JobSummaries) == 0 {
-		return nil, nil, errNoJobForBranch
+		return nil, errNoJobForBranch
 	}
 
-	latestJob = &resp.JobSummaries[0]
-	if branch.ActiveJobId != nil {
+	if includeLatest {
+		jobSummaries = append(jobSummaries, resp.JobSummaries[0])
+	}
+
+	for _, id := range ids {
+		wantJobID, _ := strconv.Atoi(id)
 		for _, j := range resp.JobSummaries {
 			jobID, _ := strconv.Atoi(*j.JobId)
-			activeJobID, _ := strconv.Atoi(*branch.ActiveJobId)
-			if jobID == activeJobID {
-				activeJob = &j
+			if jobID == wantJobID {
+				jobSummaries = append(jobSummaries, j)
+				break
 			}
 		}
+
+	}
+
+	return jobSummaries, nil
+}
+
+func (amp *AmplifyPreview) GetLatestAndActiveJobs(ctx context.Context, branch *types.Branch) (latestJob, activeJob *types.JobSummary, err error) {
+	var jobIDs []string
+	if branch.ActiveJobId != nil {
+		jobIDs = append(jobIDs, *branch.ActiveJobId)
+	}
+	jobSummaries, err := amp.findJobsByID(ctx, branch, true, jobIDs...)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if len(jobSummaries) > 0 {
+		latestJob = &jobSummaries[0]
+	}
+	if len(jobSummaries) > 1 {
+		activeJob = &jobSummaries[1]
 	}
 
 	return latestJob, activeJob, nil
 }
 
+func (amp *AmplifyPreview) WaitForJobCompletion(ctx context.Context, branch *types.Branch, job *types.JobSummary) (currentJob, activeJob *types.JobSummary, err error) {
+	for i := 0; i < jobWaitTimeAttempts; i++ {
+		jobSummaries, err := amp.findJobsByID(ctx, branch, true, *job.JobId)
+		if err != nil {
+			return nil, nil, err
+		}
+		if len(jobSummaries) > 0 {
+			currentJob = &jobSummaries[0]
+		}
+		if len(jobSummaries) > 1 {
+			activeJob = &jobSummaries[1]
+		}
+		if isAmplifyJobCompleted(currentJob) {
+			break
+		}
+
+		logger.Info("Job is not in a completed state yet. Sleeping...",
+			logKeyBranchName, amp.branchName, "job_status", currentJob.Status, "job_id", *currentJob.JobId, "attempts_left", jobWaitTimeAttempts-i)
+		time.Sleep(jobWaitSleepTime)
+	}
+
+	return currentJob, activeJob, nil
+}
+
 func appIDFromBranchARN(branchArn string) (string, error) {
 	parsedArn, err := arn.Parse(branchArn)
 	if err != nil {

tools/amplify-preview/github.go

@@ -25,7 +25,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/gravitational/shared-workflows/libs/github"
 )
 
@@ -51,7 +50,7 @@ func postPreviewURL(ctx context.Context, commentBody string) error {
 	}
 
 	targetComment := github.CommentTraits{
-		BodyContains: aws.String(amplifyMarkdownHeader),
+		BodyContains: amplifyMarkdownHeader,
 	}
 
 	currentPR := github.IssueIdentifier{

tools/amplify-preview/main.go

@@ -19,6 +19,7 @@ package main
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log/slog"
 	"os"
 	"strings"
@@ -52,85 +53,98 @@ func main() {
 	kingpin.Parse()
 	ctx := context.Background()
 
-	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRetryer(func() aws.Retryer {
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRetryer(func() aws.Retryer {
 		return retry.AddWithMaxAttempts(retry.NewStandard(), 10)
 	}))
 	if err != nil {
 		logger.Error("failed to load configuration", "error", err)
 		os.Exit(1)
 	}
 
+	if amplifyAppIDs == nil || len(*amplifyAppIDs) == 0 {
+		logger.Error("expected one more more Amplify App IDs")
+		os.Exit(1)
+	}
+
 	amp := AmplifyPreview{
-		client: amplify.NewFromConfig(cfg),
-		appIDs: func() []string {
-			if len(*amplifyAppIDs) == 1 {
-				// kingpin env variables are separated by new lines, and there is no way to change the behavior
-				// https://github.com/alecthomas/kingpin/issues/249
-				return strings.Split((*amplifyAppIDs)[0], ",")
-			}
-			return *amplifyAppIDs
-		}(),
+		client:     amplify.NewFromConfig(cfg),
+		branchName: *gitBranchName,
+		appIDs:     *amplifyAppIDs,
 	}
 
-	// Check if Amplify branch is already connected to one of the Amplify Apps
-	branch, err := amp.FindExistingBranch(ctx, *gitBranchName)
-	if err != nil {
-		if !*createBranches && !errors.Is(err, errNoJobForBranch) {
-			logger.Error("failed to lookup branch", logKeyBranchName, *gitBranchName, "error", err)
-			os.Exit(1)
-		}
+	if len(amp.appIDs) == 1 {
+		// kingpin env variables are separated by new lines, and there is no way to change the behavior
+		// https://github.com/alecthomas/kingpin/issues/249
+		amp.appIDs = strings.Split(amp.appIDs[0], ",")
+	}
 
-		// If branch wasn't found, and branch creation enabled - create new branch
-		branch, err = amp.CreateBranch(ctx, *gitBranchName)
-		if err != nil {
-			logger.Error("failed to create branch", logKeyBranchName, *gitBranchName, "error", err)
-			os.Exit(1)
-		}
+	if err := execute(ctx, amp); err != nil {
+		logger.Error("execution failed", "error", err)
+		os.Exit(1)
 	}
+}
 
-	// check if existing branch was/being deployed already
-	latestJob, activeJob, err := amp.GetLatestAndActiveJobs(ctx, branch)
+func execute(ctx context.Context, amp AmplifyPreview) error {
+	branch, err := ensureAmplifyBranch(ctx, amp)
 	if err != nil {
-		if !*createBranches && !errors.Is(err, errNoJobForBranch) {
-			logger.Error("failed to get amplify job", logKeyBranchName, *gitBranchName, "error", err)
-			os.Exit(1)
-		}
+		return err
+	}
 
-		// if job not found and branch was just created - start new job
-		latestJob, err = amp.StartJob(ctx, branch)
-		if err != nil {
-			logger.Error("failed to start amplify job", logKeyBranchName, *gitBranchName, "error", err)
-			os.Exit(1)
-		}
+	currentJob, activeJob, err := ensureAmplifyDeployment(ctx, amp, branch)
+	if err != nil {
+		return err
 	}
 
-	if err := postPreviewURL(ctx, amplifyJobsToMarkdown(branch, latestJob, activeJob)); err != nil {
-		logger.Error("failed to post preview URL", "error", err)
-		os.Exit(1)
+	if err := postPreviewURL(ctx, amplifyJobsToMarkdown(branch, currentJob, activeJob)); err != nil {
+		return fmt.Errorf("failed to post preview URL: %w", err)
 	}
+
 	logger.Info("Successfully posted PR comment")
 
 	if *wait {
-		for i := 0; !isAmplifyJobCompleted(latestJob) && i < jobWaitTimeAttempts; i++ {
-			latestJob, activeJob, err = amp.GetLatestAndActiveJobs(ctx, branch)
-			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", latestJob.Status, "job_id", *latestJob.JobId, "attempts_left", jobWaitTimeAttempts-i)
-			time.Sleep(jobWaitSleepTime)
+		currentJob, activeJob, err = amp.WaitForJobCompletion(ctx, branch, currentJob)
+		if err != nil {
+			return fmt.Errorf("failed to follow job status: %w", err)
 		}
 
-		if err := postPreviewURL(ctx, amplifyJobsToMarkdown(branch, latestJob, activeJob)); err != nil {
-			logger.Error("failed to post preview URL", "error", err)
-			os.Exit(1)
+		// update comment when job is completed
+		if err := postPreviewURL(ctx, amplifyJobsToMarkdown(branch, currentJob, activeJob)); err != nil {
+			return fmt.Errorf("failed to post preview URL: %w", err)
 		}
 	}
 
-	if latestJob.Status == types.JobStatusFailed {
-		logger.Error("amplify job is in failed state", logKeyBranchName, *gitBranchName, "job_status", latestJob.Status, "job_id", *latestJob.JobId)
-		os.Exit(1)
+	if currentJob.Status == types.JobStatusFailed {
+		logger.Error("amplify job is in failed state", logKeyBranchName, amp.branchName, "job_status", currentJob.Status, "job_id", *currentJob.JobId)
+		return fmt.Errorf("amplify job is in %q state", currentJob.Status)
+	}
+
+	return nil
+}
+
+// ensureAmplifyBranch checks if git branch is connected to amplify across multiple amplify apps
+// if "create-branch" is enabled, then branch is created if not found, otherwise returns error
+func ensureAmplifyBranch(ctx context.Context, amp AmplifyPreview) (*types.Branch, error) {
+	branch, err := amp.FindExistingBranch(ctx)
+	if err == nil {
+		return branch, nil
+	} else if errors.Is(err, errBranchNotFound) && *createBranches {
+		return amp.CreateBranch(ctx)
+	} else {
+		return nil, fmt.Errorf("failed to lookup branch %q: %w", amp.branchName, err)
+	}
+}
+
+// ensureAmplifyDeployment lists deployments and checks for latest and active (the one that's currently live) deployments
+// if this branch has no deployments yet and "create-branch" is enabled, then deployment will be kicked off
+// this is because when new branch is created on Amplify and "AutoBuild" is enabled, it won't start the deployment until next commit
+func ensureAmplifyDeployment(ctx context.Context, amp AmplifyPreview, branch *types.Branch) (currentJob, activeJob *types.JobSummary, err error) {
+	currentJob, activeJob, err = amp.GetLatestAndActiveJobs(ctx, branch)
+	if err == nil {
+		return currentJob, activeJob, nil
+	} else if errors.Is(err, errNoJobForBranch) && *createBranches {
+		currentJob, err = amp.StartJob(ctx, branch)
+		return currentJob, activeJob, err
+	} else {
+		return nil, nil, fmt.Errorf("failed to lookup amplify job for branch %q: %w", amp.branchName, err)
 	}
 }