From 73340da32226ed5b8aa4234377caa3e90bb76d4e Mon Sep 17 00:00:00 2001 From: Gabriel Adrian Samfira Date: Mon, 28 Apr 2025 13:37:33 +0000 Subject: [PATCH] Add RateLimit() function to gh client Signed-off-by: Gabriel Adrian Samfira --- runner/common/util.go | 1 + runner/pool/stub_client.go | 4 ++ test/integration/organizations_test.go | 4 +- test/integration/repositories_test.go | 4 +- util/github/client.go | 70 +++++++++++++++++--------- 5 files changed, 54 insertions(+), 29 deletions(-) diff --git a/runner/common/util.go b/runner/common/util.go index 71b1849f..06751aa3 100644 --- a/runner/common/util.go +++ b/runner/common/util.go @@ -18,6 +18,7 @@ type GithubEntityOperations interface { ListEntityRunners(ctx context.Context, opts *github.ListRunnersOptions) (*github.Runners, *github.Response, error) ListEntityRunnerApplicationDownloads(ctx context.Context) ([]*github.RunnerApplicationDownload, *github.Response, error) RemoveEntityRunner(ctx context.Context, runnerID int64) error + RateLimit(ctx context.Context) (*github.RateLimits, error) CreateEntityRegistrationToken(ctx context.Context) (*github.RegistrationToken, *github.Response, error) GetEntityJITConfig(ctx context.Context, instance string, pool params.Pool, labels []string) (jitConfigMap map[string]string, runner *github.Runner, err error) diff --git a/runner/pool/stub_client.go b/runner/pool/stub_client.go index d01c834e..2518ce9c 100644 --- a/runner/pool/stub_client.go +++ b/runner/pool/stub_client.go @@ -64,3 +64,7 @@ func (s *stubGithubClient) GetEntity() params.GithubEntity { func (s *stubGithubClient) GithubBaseURL() *url.URL { return nil } + +func (s *stubGithubClient) RateLimit(_ context.Context) (*github.RateLimits, error) { + return nil, s.err +} diff --git a/test/integration/organizations_test.go b/test/integration/organizations_test.go index 0151d2fc..a96e625c 100644 --- a/test/integration/organizations_test.go +++ b/test/integration/organizations_test.go @@ -99,8 +99,8 @@ func getGhOrgWebhook(url, ghToken, orgName string) (*github.Hook, error) { } for _, hook := range ghOrgHooks { - hookURL, ok := hook.Config["url"].(string) - if ok && hookURL == url { + hookURL := hook.Config.GetURL() + if hookURL == url { return hook, nil } } diff --git a/test/integration/repositories_test.go b/test/integration/repositories_test.go index bcf948e5..7b396ffc 100644 --- a/test/integration/repositories_test.go +++ b/test/integration/repositories_test.go @@ -109,8 +109,8 @@ func getGhRepoWebhook(url, ghToken, orgName, repoName string) (*github.Hook, err } for _, hook := range ghRepoHooks { - hookURL, ok := hook.Config["url"].(string) - if ok && hookURL == url { + hookURL := hook.Config.GetURL() + if hookURL == url { return hook, nil } } diff --git a/util/github/client.go b/util/github/client.go index 1b899913..fcd661fa 100644 --- a/util/github/client.go +++ b/util/github/client.go @@ -37,6 +37,7 @@ type githubClient struct { org *github.OrganizationsService repo *github.RepositoriesService enterprise *github.EnterpriseService + rateLimit *github.RateLimitService entity params.GithubEntity cli *github.Client @@ -226,6 +227,38 @@ func (g *githubClient) ListEntityRunnerApplicationDownloads(ctx context.Context) return ret, response, err } +func parseError(response *github.Response, err error) error { + switch response.StatusCode { + case http.StatusNotFound: + return runnerErrors.ErrNotFound + case http.StatusUnauthorized: + return runnerErrors.ErrUnauthorized + case http.StatusUnprocessableEntity: + return runnerErrors.ErrBadRequest + default: + if response.StatusCode >= 100 && response.StatusCode < 300 { + return nil + } + if err != nil { + errResp := &github.ErrorResponse{} + if errors.As(err, &errResp) && errResp.Response != nil { + switch errResp.Response.StatusCode { + case http.StatusNotFound: + return runnerErrors.ErrNotFound + case http.StatusUnauthorized: + return runnerErrors.ErrUnauthorized + case http.StatusUnprocessableEntity: + return runnerErrors.ErrBadRequest + default: + return err + } + } + return err + } + return errors.New("unknown error") + } +} + func (g *githubClient) RemoveEntityRunner(ctx context.Context, runnerID int64) error { var response *github.Response var err error @@ -254,30 +287,8 @@ func (g *githubClient) RemoveEntityRunner(ctx context.Context, runnerID int64) e return errors.New("invalid entity type") } - switch response.StatusCode { - case http.StatusNotFound: - return runnerErrors.NewNotFoundError("runner %d not found", runnerID) - case http.StatusUnauthorized: - return runnerErrors.ErrUnauthorized - case http.StatusUnprocessableEntity: - return runnerErrors.NewBadRequestError("cannot remove runner %d in its current state", runnerID) - default: - if err != nil { - errResp := &github.ErrorResponse{} - if errors.As(err, &errResp) && errResp.Response != nil { - switch errResp.Response.StatusCode { - case http.StatusNotFound: - return runnerErrors.NewNotFoundError("runner %d not found", runnerID) - case http.StatusUnauthorized: - return runnerErrors.ErrUnauthorized - case http.StatusUnprocessableEntity: - return runnerErrors.NewBadRequestError("cannot remove runner %d in its current state", runnerID) - default: - return errors.Wrap(err, "removing runner") - } - } - return errors.Wrap(err, "removing runner") - } + if err := parseError(response, err); err != nil { + return errors.Wrapf(err, "removing runner %d", runnerID) } return nil @@ -411,7 +422,7 @@ func (g *githubClient) GetEntityJITConfig(ctx context.Context, instance string, Labels: labels, // nolint:golangci-lint,godox // TODO(gabriel-samfira): Should we make this configurable? - WorkFolder: github.String("_work"), + WorkFolder: github.Ptr("_work"), } metrics.GithubOperationCount.WithLabelValues( @@ -463,6 +474,14 @@ func (g *githubClient) GetEntityJITConfig(ctx context.Context, instance string, return jitConfig, ret.Runner, nil } +func (g *githubClient) RateLimit(ctx context.Context) (*github.RateLimits, error) { + limits, resp, err := g.rateLimit.Get(ctx) + if err := parseError(resp, err); err != nil { + return nil, fmt.Errorf("getting rate limit: %w", err) + } + return limits, nil +} + func (g *githubClient) GetEntity() params.GithubEntity { return g.entity } @@ -494,6 +513,7 @@ func Client(ctx context.Context, entity params.GithubEntity) (common.GithubClien org: ghClient.Organizations, repo: ghClient.Repositories, enterprise: ghClient.Enterprise, + rateLimit: ghClient.RateLimit, cli: ghClient, entity: entity, }