Add pull request support to SCM generator

This commit adds pull request support to SCM generator so the generator can create ArgoCD apps for PRs as well. Fixes argoproj#466 Signed-off-by: Fardin Khanjani <[email protected]>
fardin01 · Jan 20, 2022 · e024751 · e024751
1 parent c8fb714
commit e024751
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 13 deletions.
diff --git a/api/v1alpha1/applicationset_types.go b/api/v1alpha1/applicationset_types.go
@@ -269,6 +269,8 @@ type SCMProviderGeneratorGithub struct {
 	TokenRef *SecretRef `json:"tokenRef,omitempty"`
 	// Scan all branches instead of just the default branch.
 	AllBranches bool `json:"allBranches,omitempty"`
+	// Scan all pull requests
+	AllPullRequests bool `json:"allPullRequests,omitempty"`
 }
 
 // SCMProviderGeneratorGitlab defines a connection info specific to Gitlab.

diff --git a/applicationset b/applicationset
diff --git a/pkg/generators/scm_provider.go b/pkg/generators/scm_provider.go
@@ -64,7 +64,7 @@ func (g *SCMProviderGenerator) GenerateParams(appSetGenerator *argoprojiov1alpha
 		if err != nil {
 			return nil, fmt.Errorf("error fetching Github token: %v", err)
 		}
-		provider, err = scm_provider.NewGithubProvider(ctx, providerConfig.Github.Organization, token, providerConfig.Github.API, providerConfig.Github.AllBranches)
+		provider, err = scm_provider.NewGithubProvider(ctx, providerConfig.Github.Organization, token, providerConfig.Github.API, providerConfig.Github.AllBranches, providerConfig.Github.AllPullRequests)
 		if err != nil {
 			return nil, fmt.Errorf("error initializing Github service: %v", err)
 		}
@@ -92,7 +92,7 @@ func (g *SCMProviderGenerator) GenerateParams(appSetGenerator *argoprojiov1alpha
 			"organization": repo.Organization,
 			"repository":   repo.Repository,
 			"url":          repo.URL,
-			"branch":       repo.Branch,
+			"revision":     repo.Revision,
 			"sha":          repo.SHA,
 			"labels":       strings.Join(repo.Labels, ","),
 		})

diff --git a/pkg/services/scm_provider/github.go b/pkg/services/scm_provider/github.go
@@ -4,20 +4,22 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strconv"
 
 	"github.com/google/go-github/v35/github"
 	"golang.org/x/oauth2"
 )
 
 type GithubProvider struct {
-	client       *github.Client
-	organization string
-	allBranches  bool
+	client          *github.Client
+	organization    string
+	allBranches     bool
+	allPullRequests bool
 }
 
 var _ SCMProviderService = &GithubProvider{}
 
-func NewGithubProvider(ctx context.Context, organization string, token string, url string, allBranches bool) (*GithubProvider, error) {
+func NewGithubProvider(ctx context.Context, organization string, token string, url string, allBranches bool, allPullRequests bool) (*GithubProvider, error) {
 	var ts oauth2.TokenSource
 	// Undocumented environment variable to set a default token, to be used in testing to dodge anonymous rate limits.
 	if token == "" {
@@ -39,7 +41,8 @@ func NewGithubProvider(ctx context.Context, organization string, token string, u
 			return nil, err
 		}
 	}
-	return &GithubProvider{client: client, organization: organization, allBranches: allBranches}, nil
+
+	return &GithubProvider{client: client, organization: organization, allBranches: allBranches, allPullRequests: allPullRequests}, nil
 }
 
 func (g *GithubProvider) ListRepos(ctx context.Context, cloneProtocol string) ([]*Repository, error) {
@@ -74,11 +77,29 @@ func (g *GithubProvider) ListRepos(ctx context.Context, cloneProtocol string) ([
 					Organization: githubRepo.Owner.GetLogin(),
 					Repository:   githubRepo.GetName(),
 					URL:          url,
-					Branch:       branch.GetName(),
+					Revision:     branch.GetName(),
 					SHA:          branch.GetCommit().GetSHA(),
 					Labels:       githubRepo.Topics,
 				})
 			}
+
+			if g.allPullRequests {
+				pullRequests, err := g.listPullRequests(ctx, githubRepo)
+				if err != nil {
+					return nil, fmt.Errorf("error listing pull requests for %s/%s: %v", githubRepo.Owner.GetLogin(), githubRepo.GetName(), err)
+				}
+
+				for _, pr := range pullRequests {
+					repos = append(repos, &Repository{
+						Organization: githubRepo.Owner.GetLogin(),
+						Repository:   githubRepo.GetName(),
+						URL:          url,
+						Revision:     strconv.FormatInt(int64(pr.GetNumber()), 10), // PR number is an int
+						SHA:          pr.GetHead().GetSHA(),
+						Labels:       githubRepo.Topics,
+					})
+				}
+			}
 		}
 		if resp.NextPage == 0 {
 			break
@@ -90,7 +111,7 @@ func (g *GithubProvider) ListRepos(ctx context.Context, cloneProtocol string) ([
 
 func (g *GithubProvider) RepoHasPath(ctx context.Context, repo *Repository, path string) (bool, error) {
 	_, _, resp, err := g.client.Repositories.GetContents(ctx, repo.Organization, repo.Repository, path, &github.RepositoryContentGetOptions{
-		Ref: repo.Branch,
+		Ref: repo.Revision,
 	})
 	// 404s are not an error here, just a normal false.
 	if resp != nil && resp.StatusCode == 404 {
@@ -132,3 +153,29 @@ func (g *GithubProvider) listBranches(ctx context.Context, repo *github.Reposito
 	}
 	return branches, nil
 }
+
+func (g *GithubProvider) listPullRequests(ctx context.Context, repo *github.Repository) ([]github.PullRequest, error) {
+	opt := &github.PullRequestListOptions{
+		ListOptions: github.ListOptions{PerPage: 100},
+	}
+
+	pullRequests := []github.PullRequest{}
+
+	for {
+		allPullRequests, resp, err := g.client.PullRequests.List(ctx, repo.Owner.GetLogin(), repo.GetName(), opt)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, pr := range allPullRequests {
+			pullRequests = append(pullRequests, *pr)
+		}
+
+		if resp.NextPage == 0 {
+			break
+		}
+		opt.Page = resp.NextPage
+	}
+
+	return pullRequests, nil
+}
diff --git a/pkg/services/scm_provider/gitlab.go b/pkg/services/scm_provider/gitlab.go
@@ -72,7 +72,7 @@ func (g *GitlabProvider) ListRepos(ctx context.Context, cloneProtocol string) ([
 					Organization: gitlabRepo.Namespace.FullPath,
 					Repository:   gitlabRepo.Path,
 					URL:          url,
-					Branch:       branch.Name,
+					Revision:     branch.Name,
 					SHA:          branch.Commit.ID,
 					Labels:       gitlabRepo.TagList,
 				})
@@ -93,7 +93,7 @@ func (g *GitlabProvider) RepoHasPath(_ context.Context, repo *Repository, path s
 	}
 	_, resp, err := g.client.Repositories.ListTree(p.ID, &gitlab.ListTreeOptions{
 		Path: &path,
-		Ref:  &repo.Branch,
+		Ref:  &repo.Revision,
 	})
 	if err != nil {
 		return false, err

diff --git a/pkg/services/scm_provider/types.go b/pkg/services/scm_provider/types.go
@@ -10,7 +10,7 @@ type Repository struct {
 	Organization string
 	Repository   string
 	URL          string
-	Branch       string
+	Revision     string // Would love suggestions on better alternatives to "Revision". Something that can represent both a branch and a PR
 	SHA          string
 	Labels       []string
 }

diff --git a/pkg/services/scm_provider/utils.go b/pkg/services/scm_provider/utils.go
@@ -44,7 +44,7 @@ func matchFilter(ctx context.Context, provider SCMProviderService, repo *Reposit
 		return false, nil
 	}
 
-	if filter.BranchMatch != nil && !filter.BranchMatch.MatchString(repo.Branch) {
+	if filter.BranchMatch != nil && !filter.BranchMatch.MatchString(repo.Revision) {
 		return false, nil
 	}