Skip to content

Commit e048b88

Browse files
lunnylunny
committed
multiple tokens support for migrating from github
1 parent 5636204 commit e048b88

File tree

1 file changed

+92
-57
lines changed

1 file changed

+92
-57
lines changed

‎modules/migrations/github.go‎

Lines changed: 92 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -68,14 +68,15 @@ func (f *GithubDownloaderV3Factory) GitServiceType() structs.GitServiceType {
6868
// from github via APIv3
6969
type GithubDownloaderV3 struct {
7070
base.NullDownloader
71-
ctx context.Context
72-
client *github.Client
73-
repoOwner string
74-
repoName string
75-
userName string
76-
password string
77-
rate *github.Rate
78-
maxPerPage int
71+
ctx context.Context
72+
clients []*github.Client
73+
repoOwner string
74+
repoName string
75+
userName string
76+
password string
77+
rates []*github.Rate
78+
curClientIdx int
79+
maxPerPage int
7980
}
8081

8182
// NewGithubDownloaderV3 creates a github Downloader via github v3 API
@@ -97,15 +98,31 @@ func NewGithubDownloaderV3(ctx context.Context, baseURL, userName, password, tok
9798
},
9899
},
99100
}
101+
100102
if token != "" {
101-
ts := oauth2.StaticTokenSource(
102-
&oauth2.Token{AccessToken: token},
103-
)
104-
client = oauth2.NewClient(downloader.ctx, ts)
105-
}
106-
downloader.client = github.NewClient(client)
107-
if baseURL != "https://github.com" {
108-
downloader.client, _ = github.NewEnterpriseClient(baseURL, baseURL, client)
103+
tokens := strings.Split(token, ",")
104+
for _, token := range tokens {
105+
token = strings.TrimSpace(token)
106+
107+
ts := oauth2.StaticTokenSource(
108+
&oauth2.Token{AccessToken: token},
109+
)
110+
client = oauth2.NewClient(downloader.ctx, ts)
111+
112+
githubClient := github.NewClient(client)
113+
if baseURL != "https://github.com" {
114+
githubClient, _ = github.NewEnterpriseClient(baseURL, baseURL, client)
115+
}
116+
downloader.clients = append(downloader.clients, githubClient)
117+
downloader.rates = append(downloader.rates, nil)
118+
}
119+
} else {
120+
githubClient := github.NewClient(client)
121+
if baseURL != "https://github.com" {
122+
githubClient, _ = github.NewEnterpriseClient(baseURL, baseURL, client)
123+
}
124+
downloader.clients = append(downloader.clients, githubClient)
125+
downloader.rates = append(downloader.rates, nil)
109126
}
110127
return &downloader
111128
}
@@ -116,8 +133,18 @@ func (g *GithubDownloaderV3) SetContext(ctx context.Context) {
116133
}
117134

118135
func (g *GithubDownloaderV3) sleep() {
119-
for g.rate != nil && g.rate.Remaining <= GithubLimitRateRemaining {
120-
timer := time.NewTimer(time.Until(g.rate.Reset.Time))
136+
var recentIdx int
137+
var maxRemaining int
138+
for i := 0; i < len(g.clients); i++ {
139+
if g.rates[i] != nil && g.rates[i].Remaining > maxRemaining {
140+
maxRemaining = g.rates[i].Remaining
141+
recentIdx = i
142+
}
143+
}
144+
g.curClientIdx = recentIdx
145+
146+
for g.rates[g.curClientIdx] != nil && g.rates[g.curClientIdx].Remaining <= GithubLimitRateRemaining {
147+
timer := time.NewTimer(time.Until(g.rates[g.curClientIdx].Reset.Time))
121148
select {
122149
case <-g.ctx.Done():
123150
util.StopTimer(timer)
@@ -127,35 +154,43 @@ func (g *GithubDownloaderV3) sleep() {
127154

128155
err := g.RefreshRate()
129156
if err != nil {
130-
log.Error("g.client.RateLimits: %s", err)
157+
log.Error("g.getClient().RateLimits: %s", err)
131158
}
132159
}
133160
}
134161

135162
// RefreshRate update the current rate (doesn't count in rate limit)
136163
func (g *GithubDownloaderV3) RefreshRate() error {
137-
rates, _, err := g.client.RateLimits(g.ctx)
164+
rates, _, err := g.getClient().RateLimits(g.ctx)
138165
if err != nil {
139166
// if rate limit is not enabled, ignore it
140167
if strings.Contains(err.Error(), "404") {
141-
g.rate = nil
168+
g.setRate(nil)
142169
return nil
143170
}
144171
return err
145172
}
146173

147-
g.rate = rates.GetCore()
174+
g.setRate(rates.GetCore())
148175
return nil
149176
}
150177

178+
func (g *GithubDownloaderV3) getClient() *github.Client {
179+
return g.clients[g.curClientIdx]
180+
}
181+
182+
func (g *GithubDownloaderV3) setRate(rate *github.Rate) {
183+
g.rates[g.curClientIdx] = rate
184+
}
185+
151186
// GetRepoInfo returns a repository information
152187
func (g *GithubDownloaderV3) GetRepoInfo() (*base.Repository, error) {
153188
g.sleep()
154-
gr, resp, err := g.client.Repositories.Get(g.ctx, g.repoOwner, g.repoName)
189+
gr, resp, err := g.getClient().Repositories.Get(g.ctx, g.repoOwner, g.repoName)
155190
if err != nil {
156191
return nil, err
157192
}
158-
g.rate = &resp.Rate
193+
g.setRate(&resp.Rate)
159194

160195
// convert github repo to stand Repo
161196
return &base.Repository{
@@ -172,11 +207,11 @@ func (g *GithubDownloaderV3) GetRepoInfo() (*base.Repository, error) {
172207
// GetTopics return github topics
173208
func (g *GithubDownloaderV3) GetTopics() ([]string, error) {
174209
g.sleep()
175-
r, resp, err := g.client.Repositories.Get(g.ctx, g.repoOwner, g.repoName)
210+
r, resp, err := g.getClient().Repositories.Get(g.ctx, g.repoOwner, g.repoName)
176211
if err != nil {
177212
return nil, err
178213
}
179-
g.rate = &resp.Rate
214+
g.setRate(&resp.Rate)
180215
return r.Topics, nil
181216
}
182217

@@ -186,7 +221,7 @@ func (g *GithubDownloaderV3) GetMilestones() ([]*base.Milestone, error) {
186221
var milestones = make([]*base.Milestone, 0, perPage)
187222
for i := 1; ; i++ {
188223
g.sleep()
189-
ms, resp, err := g.client.Issues.ListMilestones(g.ctx, g.repoOwner, g.repoName,
224+
ms, resp, err := g.getClient().Issues.ListMilestones(g.ctx, g.repoOwner, g.repoName,
190225
&github.MilestoneListOptions{
191226
State: "all",
192227
ListOptions: github.ListOptions{
@@ -196,7 +231,7 @@ func (g *GithubDownloaderV3) GetMilestones() ([]*base.Milestone, error) {
196231
if err != nil {
197232
return nil, err
198233
}
199-
g.rate = &resp.Rate
234+
g.setRate(&resp.Rate)
200235

201236
for _, m := range ms {
202237
var state = "open"
@@ -234,15 +269,15 @@ func (g *GithubDownloaderV3) GetLabels() ([]*base.Label, error) {
234269
var labels = make([]*base.Label, 0, perPage)
235270
for i := 1; ; i++ {
236271
g.sleep()
237-
ls, resp, err := g.client.Issues.ListLabels(g.ctx, g.repoOwner, g.repoName,
272+
ls, resp, err := g.getClient().Issues.ListLabels(g.ctx, g.repoOwner, g.repoName,
238273
&github.ListOptions{
239274
Page: i,
240275
PerPage: perPage,
241276
})
242277
if err != nil {
243278
return nil, err
244279
}
245-
g.rate = &resp.Rate
280+
g.setRate(&resp.Rate)
246281

247282
for _, label := range ls {
248283
labels = append(labels, convertGithubLabel(label))
@@ -291,12 +326,12 @@ func (g *GithubDownloaderV3) convertGithubRelease(rel *github.RepositoryRelease)
291326
Updated: asset.UpdatedAt.Time,
292327
DownloadFunc: func() (io.ReadCloser, error) {
293328
g.sleep()
294-
asset, redirectURL, err := g.client.Repositories.DownloadReleaseAsset(g.ctx, g.repoOwner, g.repoName, assetID, nil)
329+
asset, redirectURL, err := g.getClient().Repositories.DownloadReleaseAsset(g.ctx, g.repoOwner, g.repoName, assetID, nil)
295330
if err != nil {
296331
return nil, err
297332
}
298333
if err := g.RefreshRate(); err != nil {
299-
log.Error("g.client.RateLimits: %s", err)
334+
log.Error("g.getClient().RateLimits: %s", err)
300335
}
301336
if asset == nil {
302337
if redirectURL != "" {
@@ -308,7 +343,7 @@ func (g *GithubDownloaderV3) convertGithubRelease(rel *github.RepositoryRelease)
308343
resp, err := httpClient.Do(req)
309344
err1 := g.RefreshRate()
310345
if err1 != nil {
311-
log.Error("g.client.RateLimits: %s", err1)
346+
log.Error("g.getClient().RateLimits: %s", err1)
312347
}
313348
if err != nil {
314349
return nil, err
@@ -330,15 +365,15 @@ func (g *GithubDownloaderV3) GetReleases() ([]*base.Release, error) {
330365
var releases = make([]*base.Release, 0, perPage)
331366
for i := 1; ; i++ {
332367
g.sleep()
333-
ls, resp, err := g.client.Repositories.ListReleases(g.ctx, g.repoOwner, g.repoName,
368+
ls, resp, err := g.getClient().Repositories.ListReleases(g.ctx, g.repoOwner, g.repoName,
334369
&github.ListOptions{
335370
Page: i,
336371
PerPage: perPage,
337372
})
338373
if err != nil {
339374
return nil, err
340375
}
341-
g.rate = &resp.Rate
376+
g.setRate(&resp.Rate)
342377

343378
for _, release := range ls {
344379
releases = append(releases, g.convertGithubRelease(release))
@@ -367,12 +402,12 @@ func (g *GithubDownloaderV3) GetIssues(page, perPage int) ([]*base.Issue, bool,
367402

368403
var allIssues = make([]*base.Issue, 0, perPage)
369404
g.sleep()
370-
issues, resp, err := g.client.Issues.ListByRepo(g.ctx, g.repoOwner, g.repoName, opt)
405+
issues, resp, err := g.getClient().Issues.ListByRepo(g.ctx, g.repoOwner, g.repoName, opt)
371406
if err != nil {
372407
return nil, false, fmt.Errorf("error while listing repos: %v", err)
373408
}
374409
log.Trace("Request get issues %d/%d, but in fact get %d", perPage, page, len(issues))
375-
g.rate = &resp.Rate
410+
g.setRate(&resp.Rate)
376411
for _, issue := range issues {
377412
if issue.IsPullRequest() {
378413
continue
@@ -387,14 +422,14 @@ func (g *GithubDownloaderV3) GetIssues(page, perPage int) ([]*base.Issue, bool,
387422
var reactions []*base.Reaction
388423
for i := 1; ; i++ {
389424
g.sleep()
390-
res, resp, err := g.client.Reactions.ListIssueReactions(g.ctx, g.repoOwner, g.repoName, issue.GetNumber(), &github.ListOptions{
425+
res, resp, err := g.getClient().Reactions.ListIssueReactions(g.ctx, g.repoOwner, g.repoName, issue.GetNumber(), &github.ListOptions{
391426
Page: i,
392427
PerPage: perPage,
393428
})
394429
if err != nil {
395430
return nil, false, err
396431
}
397-
g.rate = &resp.Rate
432+
g.setRate(&resp.Rate)
398433
if len(res) == 0 {
399434
break
400435
}
@@ -465,24 +500,24 @@ func (g *GithubDownloaderV3) getComments(issueContext base.IssueContext) ([]*bas
465500
}
466501
for {
467502
g.sleep()
468-
comments, resp, err := g.client.Issues.ListComments(g.ctx, g.repoOwner, g.repoName, int(issueContext.ForeignID()), opt)
503+
comments, resp, err := g.getClient().Issues.ListComments(g.ctx, g.repoOwner, g.repoName, int(issueContext.ForeignID()), opt)
469504
if err != nil {
470505
return nil, fmt.Errorf("error while listing repos: %v", err)
471506
}
472-
g.rate = &resp.Rate
507+
g.setRate(&resp.Rate)
473508
for _, comment := range comments {
474509
// get reactions
475510
var reactions []*base.Reaction
476511
for i := 1; ; i++ {
477512
g.sleep()
478-
res, resp, err := g.client.Reactions.ListIssueCommentReactions(g.ctx, g.repoOwner, g.repoName, comment.GetID(), &github.ListOptions{
513+
res, resp, err := g.getClient().Reactions.ListIssueCommentReactions(g.ctx, g.repoOwner, g.repoName, comment.GetID(), &github.ListOptions{
479514
Page: i,
480515
PerPage: g.maxPerPage,
481516
})
482517
if err != nil {
483518
return nil, err
484519
}
485-
g.rate = &resp.Rate
520+
g.setRate(&resp.Rate)
486521
if len(res) == 0 {
487522
break
488523
}
@@ -534,27 +569,27 @@ func (g *GithubDownloaderV3) GetAllComments(page, perPage int) ([]*base.Comment,
534569
}
535570

536571
g.sleep()
537-
comments, resp, err := g.client.Issues.ListComments(g.ctx, g.repoOwner, g.repoName, 0, opt)
572+
comments, resp, err := g.getClient().Issues.ListComments(g.ctx, g.repoOwner, g.repoName, 0, opt)
538573
if err != nil {
539574
return nil, false, fmt.Errorf("error while listing repos: %v", err)
540575
}
541576
var isEnd = resp.NextPage == 0
542577

543578
log.Trace("Request get comments %d/%d, but in fact get %d, next page is %d", perPage, page, len(comments), resp.NextPage)
544-
g.rate = &resp.Rate
579+
g.setRate(&resp.Rate)
545580
for _, comment := range comments {
546581
// get reactions
547582
var reactions []*base.Reaction
548583
for i := 1; ; i++ {
549584
g.sleep()
550-
res, resp, err := g.client.Reactions.ListIssueCommentReactions(g.ctx, g.repoOwner, g.repoName, comment.GetID(), &github.ListOptions{
585+
res, resp, err := g.getClient().Reactions.ListIssueCommentReactions(g.ctx, g.repoOwner, g.repoName, comment.GetID(), &github.ListOptions{
551586
Page: i,
552587
PerPage: g.maxPerPage,
553588
})
554589
if err != nil {
555590
return nil, false, err
556591
}
557-
g.rate = &resp.Rate
592+
g.setRate(&resp.Rate)
558593
if len(res) == 0 {
559594
break
560595
}
@@ -599,12 +634,12 @@ func (g *GithubDownloaderV3) GetPullRequests(page, perPage int) ([]*base.PullReq
599634
}
600635
var allPRs = make([]*base.PullRequest, 0, perPage)
601636
g.sleep()
602-
prs, resp, err := g.client.PullRequests.List(g.ctx, g.repoOwner, g.repoName, opt)
637+
prs, resp, err := g.getClient().PullRequests.List(g.ctx, g.repoOwner, g.repoName, opt)
603638
if err != nil {
604639
return nil, false, fmt.Errorf("error while listing repos: %v", err)
605640
}
606641
log.Trace("Request get pull requests %d/%d, but in fact get %d", perPage, page, len(prs))
607-
g.rate = &resp.Rate
642+
g.setRate(&resp.Rate)
608643
for _, pr := range prs {
609644
var labels = make([]*base.Label, 0, len(pr.Labels))
610645
for _, l := range pr.Labels {
@@ -615,14 +650,14 @@ func (g *GithubDownloaderV3) GetPullRequests(page, perPage int) ([]*base.PullReq
615650
var reactions []*base.Reaction
616651
for i := 1; ; i++ {
617652
g.sleep()
618-
res, resp, err := g.client.Reactions.ListIssueReactions(g.ctx, g.repoOwner, g.repoName, pr.GetNumber(), &github.ListOptions{
653+
res, resp, err := g.getClient().Reactions.ListIssueReactions(g.ctx, g.repoOwner, g.repoName, pr.GetNumber(), &github.ListOptions{
619654
Page: i,
620655
PerPage: perPage,
621656
})
622657
if err != nil {
623658
return nil, false, err
624659
}
625-
g.rate = &resp.Rate
660+
g.setRate(&resp.Rate)
626661
if len(res) == 0 {
627662
break
628663
}
@@ -693,14 +728,14 @@ func (g *GithubDownloaderV3) convertGithubReviewComments(cs []*github.PullReques
693728
var reactions []*base.Reaction
694729
for i := 1; ; i++ {
695730
g.sleep()
696-
res, resp, err := g.client.Reactions.ListPullRequestCommentReactions(g.ctx, g.repoOwner, g.repoName, c.GetID(), &github.ListOptions{
731+
res, resp, err := g.getClient().Reactions.ListPullRequestCommentReactions(g.ctx, g.repoOwner, g.repoName, c.GetID(), &github.ListOptions{
697732
Page: i,
698733
PerPage: g.maxPerPage,
699734
})
700735
if err != nil {
701736
return nil, err
702737
}
703-
g.rate = &resp.Rate
738+
g.setRate(&resp.Rate)
704739
if len(res) == 0 {
705740
break
706741
}
@@ -738,11 +773,11 @@ func (g *GithubDownloaderV3) GetReviews(context base.IssueContext) ([]*base.Revi
738773
}
739774
for {
740775
g.sleep()
741-
reviews, resp, err := g.client.PullRequests.ListReviews(g.ctx, g.repoOwner, g.repoName, int(context.ForeignID()), opt)
776+
reviews, resp, err := g.getClient().PullRequests.ListReviews(g.ctx, g.repoOwner, g.repoName, int(context.ForeignID()), opt)
742777
if err != nil {
743778
return nil, fmt.Errorf("error while listing repos: %v", err)
744779
}
745-
g.rate = &resp.Rate
780+
g.setRate(&resp.Rate)
746781
for _, review := range reviews {
747782
r := convertGithubReview(review)
748783
r.IssueIndex = context.LocalID()
@@ -752,11 +787,11 @@ func (g *GithubDownloaderV3) GetReviews(context base.IssueContext) ([]*base.Revi
752787
}
753788
for {
754789
g.sleep()
755-
reviewComments, resp, err := g.client.PullRequests.ListReviewComments(g.ctx, g.repoOwner, g.repoName, int(context.ForeignID()), review.GetID(), opt2)
790+
reviewComments, resp, err := g.getClient().PullRequests.ListReviewComments(g.ctx, g.repoOwner, g.repoName, int(context.ForeignID()), review.GetID(), opt2)
756791
if err != nil {
757792
return nil, fmt.Errorf("error while listing repos: %v", err)
758793
}
759-
g.rate = &resp.Rate
794+
g.setRate(&resp.Rate)
760795

761796
cs, err := g.convertGithubReviewComments(reviewComments)
762797
if err != nil {

0 commit comments

Comments
 (0)