From df97d6c2d308c280ba3048c4c083da9c0476c381 Mon Sep 17 00:00:00 2001
From: Amr Tamimi <amrnt0@@gmail.com>
Date: Sat, 2 Jun 2018 02:58:25 +0200
Subject: [PATCH] initial implementation to support github enterprise

---
 _site/content/posts/modules/github.md |  8 +++++-
 github/github_repo.go                 | 40 +++++++++++++++++++++------
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/_site/content/posts/modules/github.md b/_site/content/posts/modules/github.md
index 2dd3e6ca..b5a39bc0 100644
--- a/_site/content/posts/modules/github.md
+++ b/_site/content/posts/modules/github.md
@@ -26,7 +26,13 @@ wtf/github/
 ## Required ENV Variables
 
 <span class="caption">Key:</span> `WTF_GITHUB_TOKEN` <br />
-<span class="caption">Action:</span> Your <a href="https://developer.github.com/v3/oauth_authorizations/#create-a-new-authorization">Github API</a> token.
+<span class="caption">Action:</span> Your <a href="https://developer.github.com/v3/oauth_authorizations/#create-a-new-authorization">Github API</a> token or <a href="https://developer.github.com/enterprise/2.13/v3/enterprise-admin/">Github Enterprise</a> API token.
+
+<span class="caption">Key:</span> `WTF_GITHUB_BASE_URL` <br />
+<span class="caption">Action:</span> Your <a href="https://developer.github.com/enterprise/2.13/v3/enterprise-admin/">Github Enterprise</a> API URL.
+
+<span class="caption">Key:</span> `WTF_GITHUB_UPLOAD_URL` <br />
+<span class="caption">Action:</span> Your <a href="https://developer.github.com/enterprise/2.13/v3/enterprise-admin/">Github Enterprise</a> upload URL (often the same as API URL).
 
 ## Keyboard Commands
 
diff --git a/github/github_repo.go b/github/github_repo.go
index 7ced9123..532273d0 100644
--- a/github/github_repo.go
+++ b/github/github_repo.go
@@ -10,7 +10,9 @@ import (
 )
 
 type GithubRepo struct {
-	apiKey string
+	apiKey    string
+	baseURL   string
+	uploadURL string
 
 	Name         string
 	Owner        string
@@ -20,9 +22,12 @@ type GithubRepo struct {
 
 func NewGithubRepo(name, owner string) *GithubRepo {
 	repo := GithubRepo{
-		apiKey: os.Getenv("WTF_GITHUB_TOKEN"),
-		Name:   name,
-		Owner:  owner,
+		apiKey:    os.Getenv("WTF_GITHUB_TOKEN"),
+		baseURL:   os.Getenv("WTF_GITHUB_BASE_URL"),
+		uploadURL: os.Getenv("WTF_GITHUB_UPLOAD_URL"),
+
+		Name:  name,
+		Owner: owner,
 	}
 
 	return &repo
@@ -66,6 +71,19 @@ func (repo *GithubRepo) oauthClient() *http.Client {
 	return oauth2.NewClient(context.Background(), tokenService)
 }
 
+func (repo *GithubRepo) githubClient() (*ghb.Client, error) {
+	oauthClient := repo.oauthClient()
+
+	if len(repo.baseURL) > 0 {
+		if len(repo.uploadURL) == 0 {
+			repo.uploadURL = repo.baseURL
+		}
+		return ghb.NewEnterpriseClient(repo.baseURL, repo.baseURL, oauthClient)
+	}
+
+	return ghb.NewClient(oauthClient), nil
+}
+
 // myPullRequests returns a list of pull requests created by username on this repo
 func (repo *GithubRepo) myPullRequests(username string) []*ghb.PullRequest {
 	prs := []*ghb.PullRequest{}
@@ -98,8 +116,11 @@ func (repo *GithubRepo) myReviewRequests(username string) []*ghb.PullRequest {
 }
 
 func (repo *GithubRepo) loadPullRequests() ([]*ghb.PullRequest, error) {
-	oauthClient := repo.oauthClient()
-	github := ghb.NewClient(oauthClient)
+	github, err := repo.githubClient()
+
+	if err != nil {
+		return nil, err
+	}
 
 	opts := &ghb.PullRequestListOptions{}
 
@@ -113,8 +134,11 @@ func (repo *GithubRepo) loadPullRequests() ([]*ghb.PullRequest, error) {
 }
 
 func (repo *GithubRepo) loadRemoteRepository() (*ghb.Repository, error) {
-	oauthClient := repo.oauthClient()
-	github := ghb.NewClient(oauthClient)
+	github, err := repo.githubClient()
+
+	if err != nil {
+		return nil, err
+	}
 
 	repository, _, err := github.Repositories.Get(context.Background(), repo.Owner, repo.Name)