Repository: lindell/multi-gitter
Pull request #242 (2022-03-29 06:08:37)

feat(github): use graphql endpoint to get pull request status


Patch:
diff --git a/cmd/platform.go b/cmd/platform.go
index cd2a929..baceb43 100644
--- a/cmd/platform.go
+++ b/cmd/platform.go
@@ -18,7 +18,7 @@ import (
 func configurePlatform(cmd *cobra.Command) {
 	flags := cmd.Flags()
 
-	flags.StringP("base-url", "g", "", "Base URL of the (v3) GitHub API, needs to be changed if GitHub enterprise is used. Or the url to a self-hosted GitLab instance.")
+	flags.StringP("base-url", "g", "", "Base URL of the GitHub API, needs to be changed if GitHub enterprise is used. Or the url to a self-hosted GitLab instance.")
 	flags.BoolP("insecure", "", false, "Insecure controls whether a client verifies the server certificate chain and host name. Used only for Bitbucket server.")
 	flags.StringP("username", "u", "", "The Bitbucket server username.")
 	flags.StringP("token", "T", "", "The GitHub/GitLab personal access token. Can also be set using the GITHUB_TOKEN/GITLAB_TOKEN/GITEA_TOKEN/BITBUCKET_SERVER_TOKEN environment variable.")
diff --git a/internal/scm/github/github.go b/internal/scm/github/github.go
index 514ad50..2a5a9c0 100755
--- a/internal/scm/github/github.go
+++ b/internal/scm/github/github.go
@@ -49,10 +49,14 @@ func New(
 		RepositoryListing: repoListing,
 		MergeTypes:        mergeTypes,
 		token:             token,
+		baseURL:           baseURL,
 		Fork:              forkMode,
 		ForkOwner:         forkOwner,
 		SSHAuth:           sshAuth,
 		ghClient:          client,
+		httpClient: &http.Client{
+			Transport: transportMiddleware(http.DefaultTransport),
+		},
 	}, nil
 }
 
@@ -61,6 +65,7 @@ type Github struct {
 	RepositoryListing
 	MergeTypes []scm.MergeType
 	token      string
+	baseURL    string
 
 	// This determines if forks will be used when creating a prs.
 	// In this package, it mainly determines which repos are possible to make changes on
@@ -72,7 +77,8 @@ type Github struct {
 	// If set, use the SSH clone url instead of http(s)
 	SSHAuth bool
 
-	ghClient *github.Client
+	ghClient   *github.Client
+	httpClient *http.Client
 
 	// Caching of the logged in user
 	user      string
@@ -313,52 +319,106 @@ func (g *Github) addAssignees(ctx context.Context, repo repository, newPR scm.Ne
 
 // GetPullRequests gets all pull requests of with a specific branch
 func (g *Github) GetPullRequests(ctx context.Context, branchName string) ([]scm.PullRequest, error) {
-	// TODO: If this is implemented with the GitHub v4 graphql api, it would be much faster
-
 	repos, err := g.getRepositories(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	prStatuses := []scm.PullRequest{}
-	for _, r := range repos {
-		repoOwner := r.GetOwner().GetLogin()
-		repoName := r.GetName()
-		log := log.WithField("repo", fmt.Sprintf("%s/%s", repoOwner, repoName))
-
-		headOwner, err := g.headOwner(ctx, repoOwner)
-		if err != nil {
-			return nil, err
+	// The fragment is all the data needed from every repository
+	const fragment = `fragment repoProperties on Repository {
+		pullRequests(headRefName: $branchName, last: 1) {
+			nodes {
+				number
+				headRefName
+				closed
+				url
+				baseRepository {
+					name
+					owner {
+						login
+					}
+				}
+				headRepository {
+					name
+					owner {
+						login
+					}
+				}
+				commits(last: 1) {
+					nodes {
+						commit {
+							statusCheckRollup {
+								state
+							}
+						}
+					}
+				}
+			}
 		}
+	}`
+
+	// Prepare data for compiling the query.
+	// Each repository will get its own variables ($ownerX, $repoX) and be returned
+	// via and alias repoX
+	repoParameters := make([]string, len(repos))
+	repoQueries := make([]string, len(repos))
+	queryVariables := map[string]interface{}{
+		"branchName": branchName,
+	}
+	for i, repo := range repos {
+		repoParameters[i] = fmt.Sprintf("$owner%[1]d: String!, $repo%[1]d: String!", i)
+		repoQueries[i] = fmt.Sprintf("repo%[1]d: repository(owner: $owner%[1]d, name: $repo%[1]d) { ...repoProperties }", i)
+
+		queryVariables[fmt.Sprintf("owner%d", i)] = repo.GetOwner().GetLogin()
+		queryVariables[fmt.Sprintf("repo%d", i)] = repo.GetName()
+	}
+
+	// Create the final query
+	query := fmt.Sprintf(`
+		%s
+
+		query ($branchName: String!, %s) {
+			%s
+		}`,
+		fragment,
+		strings.Join(repoParameters, ", "),
+		strings.Join(repoQueries, "\n"),
+	)
 
-		log.Debug("Fetching latest pull request")
-		prs, _, err := g.ghClient.PullRequests.List(ctx, repoOwner, repoName, &github.PullRequestListOptions{
-			Head:      fmt.Sprintf("%s:%s", headOwner, branchName),
-			State:     "all",
-			Direction: "desc",
-			ListOptions: github.ListOptions{
-				PerPage: 1,
-			},
-		})
-		if err != nil {
-			return nil, err
+	result := map[string]graphqlRepo{}
+	err = g.makeGraphQLRequest(ctx, query, queryVariables, &result)
+	if err != nil {
+		return nil, err
+	}
+
+	// Fetch the repo based on name instead of looping through the map since that will
+	// guarantee the same ordering as the original repository list
+	prs := []scm.PullRequest{}
+	for i := range repos {
+		repo, ok := result[fmt.Sprintf("repo%d", i)]
+		if !ok {
+			return nil, fmt.Errorf("could not find repo%d", i)
 		}
-		if len(prs) != 1 {
+
+		if len(repo.PullRequests.Nodes) != 1 {
 			continue
 		}
-		pr := prs[0]
+		pr := repo.PullRequests.Nodes[0]
 
-		status, err := g.getPrStatus(ctx, pr)
+		// The graphql API does not have a way at query time to filter out the owner of the head branch
+		// of a PR. Therefore, we have to filter out any repo that does not match the head owner.
+		headOwner, err := g.headOwner(ctx, pr.BaseRepository.Owner.Login)
 		if err != nil {
 			return nil, err
 		}
+		if pr.HeadRepository.Owner.Login != headOwner {
+			continue
+		}
 
-		localPR := convertPullRequest(pr)
-		localPR.status = status
-		prStatuses = append(prStatuses, localPR)
+		prs = append(prs, convertGraphQLPullRequest(pr))
 	}
 
-	return prStatuses, nil
+	return prs, nil
 }
 
 func (g *Github) loggedInUser(ctx context.Context) (string, error) {
@@ -565,37 +625,6 @@ func (g *Github) GetAutocompleteRepositories(ctx context.Context, str string) ([
 	return ret, nil
 }
 
-func (g *Github) getPrStatus(ctx context.Context, pr *github.PullRequest) (scm.PullRequestStatus, error) {
-	// Determine the status of the pr
-	var status scm.PullRequestStatus
-	if pr.MergedAt != nil {
-		status = scm.PullRequestStatusMerged
-	} else if pr.ClosedAt != nil {
-		status = scm.PullRequestStatusClosed
-	} else {
-		log.Debug("Fetching the combined status of the pull request")
-		combinedStatus, _, err := g.ghClient.Repositories.GetCombinedStatus(ctx, pr.GetBase().GetUser().GetLogin(), pr.GetBase().GetRepo().GetName(), pr.GetHead().GetSHA(), nil)
-		if err != nil {
-			return scm.PullRequestStatusUnknown, err
-		}
-
-		if combinedStatus.GetTotalCount() == 0 {
-			status = scm.PullRequestStatusSuccess
-		} else {
-			switch combinedStatus.GetState() {
-			case "pending":
-				status = scm.PullRequestStatusPending
-			case "success":
-				status = scm.PullRequestStatusSuccess
-			case "failure", "error":
-				status = scm.PullRequestStatusError
-			}
-		}
-	}
-
-	return status, nil
-}
-
 // modLock is a lock that should be used whenever a modifying request is made against the GitHub API.
 // It works as a normal lock, but also makes sure that there is a buffer period of 1 second since the
 // last critical section was left. This ensures that we always wait at least one seconds between modifying requests
diff --git a/internal/scm/github/graphql.go b/internal/scm/github/graphql.go
new file mode 100644
index 0000000..700aa60
--- /dev/null
+++ b/internal/scm/github/graphql.go
@@ -0,0 +1,140 @@
+package github
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"github.com/pkg/errors"
+)
+
+func (g *Github) makeGraphQLRequest(ctx context.Context, query string, data interface{}, res interface{}) error {
+	type reqData struct {
+		Query string      `json:"query"`
+		Data  interface{} `json:"variables"`
+	}
+	rawReqData, err := json.Marshal(reqData{
+		Query: query,
+		Data:  data,
+	})
+	if err != nil {
+		return errors.WithMessage(err, "could not marshal graphql request")
+	}
+
+	graphQLURL := "https://api.github.com/graphql"
+	if g.baseURL != "" {
+		graphQLURL, err = graphQLEndpoint(g.baseURL)
+		if err != nil {
+			return errors.WithMessage(err, "could not get graphql endpoint")
+		}
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "POST", graphQLURL, bytes.NewBuffer(rawReqData))
+	if err != nil {
+		return errors.WithMessage(err, "could not create graphql request")
+	}
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", g.token))
+
+	resp, err := g.httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	resultData := struct {
+		Data   json.RawMessage `json:"data"`
+		Errors []struct {
+			Message string `json:"message"`
+		} `json:"errors"`
+	}{}
+
+	if err := json.NewDecoder(resp.Body).Decode(&resultData); err != nil {
+		return errors.WithMessage(err, "could not read graphql response body")
+	}
+
+	if len(resultData.Errors) > 0 {
+		errorsMsgs := make([]string, len(resultData.Errors))
+		for i := range resultData.Errors {
+			errorsMsgs[i] = resultData.Errors[i].Message
+		}
+		return errors.WithMessage(
+			errors.New(strings.Join(errorsMsgs, "\n")),
+			"encountered error during GraphQL query",
+		)
+	}
+
+	if err := json.Unmarshal(resultData.Data, res); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// graphQLEndpoint takes a url to a github enterprise instance (or the v3 api) and returns the url to the graphql endpoint
+func graphQLEndpoint(u string) (string, error) {
+	baseEndpoint, err := url.Parse(u)
+	if err != nil {
+		return "", err
+	}
+	if !strings.HasSuffix(baseEndpoint.Path, "/") {
+		baseEndpoint.Path += "/"
+	}
+
+	if strings.HasPrefix(baseEndpoint.Host, "api.") ||
+		strings.Contains(baseEndpoint.Host, ".api.") {
+		baseEndpoint.Path += "graphql"
+	} else {
+		baseEndpoint.Path = stripSuffixIfExist(baseEndpoint.Path, "v3/")
+		baseEndpoint.Path = stripSuffixIfExist(baseEndpoint.Path, "api/")
+		baseEndpoint.Path += "api/graphql"
+	}
+
+	return baseEndpoint.String(), nil
+}
+
+type graphqlPullRequestState string
+
+const (
+	graphqlPullRequestStateError   graphqlPullRequestState = "ERROR"
+	graphqlPullRequestStateFailure graphqlPullRequestState = "FAILURE"
+	graphqlPullRequestStatePending graphqlPullRequestState = "PENDING"
+	graphqlPullRequestStateSuccess graphqlPullRequestState = "SUCCESS"
+)
+
+type graphqlRepo struct {
+	PullRequests struct {
+		Nodes []graphqlPR `json:"nodes"`
+	} `json:"pullRequests"`
+}
+
+type graphqlPR struct {
+	Number         int    `json:"number"`
+	HeadRefName    string `json:"headRefName"`
+	Closed         bool   `json:"closed"`
+	URL            string `json:"url"`
+	BaseRepository struct {
+		Name  string `json:"name"`
+		Owner struct {
+			Login string `json:"login"`
+		} `json:"owner"`
+	} `json:"baseRepository"`
+	HeadRepository struct {
+		Name  string `json:"name"`
+		Owner struct {
+			Login string `json:"login"`
+		} `json:"owner"`
+	} `json:"headRepository"`
+	Commits struct {
+		Nodes []struct {
+			Commit struct {
+				StatusCheckRollup struct {
+					State *graphqlPullRequestState `json:"state"`
+				} `json:"statusCheckRollup"`
+			} `json:"commit"`
+		} `json:"nodes"`
+	} `json:"commits"`
+}
diff --git a/internal/scm/github/pullrequest.go b/internal/scm/github/pullrequest.go
index 9b72e1b..db88eed 100644
--- a/internal/scm/github/pullrequest.go
+++ b/internal/scm/github/pullrequest.go
@@ -20,6 +20,35 @@ func convertPullRequest(pr *github.PullRequest) pullRequest {
 	}
 }
 
+func convertGraphQLPullRequest(pr graphqlPR) pullRequest {
+	combinedStatus := pr.Commits.Nodes[0].Commit.StatusCheckRollup.State
+	status := scm.PullRequestStatusUnknown
+
+	if combinedStatus == nil {
+		status = scm.PullRequestStatusSuccess
+	} else {
+		switch *combinedStatus {
+		case graphqlPullRequestStatePending:
+			status = scm.PullRequestStatusPending
+		case graphqlPullRequestStateSuccess:
+			status = scm.PullRequestStatusSuccess
+		case graphqlPullRequestStateFailure, graphqlPullRequestStateError:
+			status = scm.PullRequestStatusError
+		}
+	}
+
+	return pullRequest{
+		ownerName:   pr.BaseRepository.Owner.Login,
+		repoName:    pr.BaseRepository.Name,
+		branchName:  pr.HeadRefName,
+		prOwnerName: pr.HeadRepository.Owner.Login,
+		prRepoName:  pr.HeadRepository.Name,
+		number:      pr.Number,
+		guiURL:      pr.URL,
+		status:      status,
+	}
+}
+
 type pullRequest struct {
 	ownerName   string
 	repoName    string
diff --git a/internal/scm/github/util.go b/internal/scm/github/util.go
index 2e926ae..9a10bb5 100644
--- a/internal/scm/github/util.go
+++ b/internal/scm/github/util.go
@@ -1,6 +1,8 @@
 package github
 
 import (
+	"strings"
+
 	"github.com/google/go-github/v39/github"
 	"github.com/lindell/multi-gitter/internal/scm"
 )
@@ -26,3 +28,10 @@ func repoMergeTypes(repo *github.Repository) []scm.MergeType {
 	}
 	return ret
 }
+
+func stripSuffixIfExist(str string, suffix string) string {
+	if strings.HasSuffix(str, suffix) {
+		return str[:len(str)-len(suffix)]
+	}
+	return str
+}


Test patch:
diff --git a/internal/scm/github/graphql_test.go b/internal/scm/github/graphql_test.go
new file mode 100644
index 0000000..ff5ed7b
--- /dev/null
+++ b/internal/scm/github/graphql_test.go
@@ -0,0 +1,27 @@
+package github
+
+import (
+	"fmt"
+	"testing"
+)
+
+func Test_graphQLEndpoint(t *testing.T) {
+	tests := []struct {
+		url  string
+		want string
+	}{
+		{url: "https://github.detsbihcs.io/api/v3", want: "https://github.detsbihcs.io/api/graphql"},
+		{url: "https://github.detsbihcs.io/api/v3/", want: "https://github.detsbihcs.io/api/graphql"},
+		{url: "https://github.detsbihcs.io/api/", want: "https://github.detsbihcs.io/api/graphql"},
+		{url: "https://github.detsbihcs.io/", want: "https://github.detsbihcs.io/api/graphql"},
+		{url: "https://api.github.detsbihcs.io/", want: "https://api.github.detsbihcs.io/graphql"},
+		{url: "https://more.api.github.detsbihcs.io/", want: "https://more.api.github.detsbihcs.io/graphql"},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("url: %s", tt.url), func(t *testing.T) {
+			if got, _ := graphQLEndpoint(tt.url); got != tt.want {
+				t.Errorf("graphQLEndpoint() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}