Merge pull request #32 from git-pkgs/add-rate-limit

Add rate limit support

Author: andrew

bitbucket/rate_limit.go

@@ -0,0 +1,12 @@
+package bitbucket
+
+import (
+	"context"
+	"fmt"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+func (f *bitbucketForge) GetRateLimit(ctx context.Context) (*forge.RateLimit, error) {
+	return nil, fmt.Errorf("getting rate limit: %w", forge.ErrNotSupported)
+}

bitbucket/rate_limit_test.go

@@ -0,0 +1,17 @@
+package bitbucket
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+func TestBitbucketRateLimitNotSupported(t *testing.T) {
+	f := New("test-token", nil)
+	_, err := f.GetRateLimit(context.Background())
+	if !errors.Is(err, forge.ErrNotSupported) {
+		t.Fatalf("expected ErrNotSupported, got %v", err)
+	}
+}

forge.go

@@ -45,6 +45,7 @@ type Forge interface {
 	Secrets() SecretService
 	Notifications() NotificationService
 	Reviews() ReviewService
+	GetRateLimit(ctx context.Context) (*RateLimit, error)
 }
 
 // Client routes requests to the appropriate Forge based on the URL domain.

forges_test.go

@@ -477,6 +477,10 @@ func (m *mockForge) Reviews() ReviewService {
 	return &mockReviewService{}
 }
 
+func (m *mockForge) GetRateLimit(_ context.Context) (*RateLimit, error) {
+	return nil, ErrNotSupported
+}
+
 type mockRepoService struct {
 	repo      *Repository
 	repos     []Repository

gitea/gitea.go

@@ -9,7 +9,10 @@ import (
 )
 
 type giteaForge struct {
-	client *gitea.Client
+	client     *gitea.Client
+	baseURL    string
+	token      string
+	httpClient *http.Client
 }
 
 // New creates a Gitea/Forgejo forge backend.
@@ -22,7 +25,7 @@ func New(baseURL, token string, hc *http.Client) forge.Forge {
 		opts = append(opts, gitea.SetHTTPClient(hc))
 	}
 	c, _ := gitea.NewClient(baseURL, opts...)
-	return &giteaForge{client: c}
+	return &giteaForge{client: c, baseURL: baseURL, token: token, httpClient: hc}
 }
 
 type giteaRepoService struct {

gitea/rate_limit.go

@@ -0,0 +1,67 @@
+package gitea
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"time"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+type giteaRateLimitResponse struct {
+	Resources struct {
+		Core struct {
+			Limit     int   `json:"limit"`
+			Remaining int   `json:"remaining"`
+			Reset     int64 `json:"reset"`
+		} `json:"core"`
+	} `json:"resources"`
+}
+
+func (f *giteaForge) GetRateLimit(ctx context.Context) (*forge.RateLimit, error) {
+	url := f.baseURL + "/api/v1/rate_limit"
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	if f.token != "" {
+		req.Header.Set("Authorization", "token "+f.token)
+	}
+
+	hc := f.httpClient
+	if hc == nil {
+		hc = http.DefaultClient
+	}
+
+	resp, err := hc.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusNotFound {
+		return nil, fmt.Errorf("getting rate limit: %w", forge.ErrNotSupported)
+	}
+	if resp.StatusCode >= 400 {
+		return nil, &forge.HTTPError{StatusCode: resp.StatusCode, URL: url}
+	}
+
+	var result giteaRateLimitResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, err
+	}
+
+	core := result.Resources.Core
+	var reset time.Time
+	if core.Reset > 0 {
+		reset = time.Unix(core.Reset, 0)
+	}
+
+	return &forge.RateLimit{
+		Limit:     core.Limit,
+		Remaining: core.Remaining,
+		Reset:     reset,
+	}, nil
+}

gitea/rate_limit_test.go

@@ -0,0 +1,60 @@
+package gitea
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+func TestGiteaGetRateLimit(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("GET /api/v1/version", giteaVersionHandler)
+	mux.HandleFunc("GET /api/v1/rate_limit", func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"resources": map[string]any{
+				"core": map[string]any{
+					"limit":     100,
+					"remaining": 98,
+					"reset":     1717243200,
+				},
+			},
+		})
+	})
+
+	srv := httptest.NewServer(mux)
+	defer srv.Close()
+
+	f := New(srv.URL, "test-token", nil)
+	rl, err := f.GetRateLimit(context.Background())
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	assertEqualInt(t, "Limit", 100, rl.Limit)
+	assertEqualInt(t, "Remaining", 98, rl.Remaining)
+	if rl.Reset.Unix() != 1717243200 {
+		t.Errorf("Reset: want unix 1717243200, got %d", rl.Reset.Unix())
+	}
+}
+
+func TestGiteaGetRateLimitNotSupported(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("GET /api/v1/version", giteaVersionHandler)
+	mux.HandleFunc("GET /api/v1/rate_limit", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusNotFound)
+	})
+
+	srv := httptest.NewServer(mux)
+	defer srv.Close()
+
+	f := New(srv.URL, "test-token", nil)
+	_, err := f.GetRateLimit(context.Background())
+	if !errors.Is(err, forge.ErrNotSupported) {
+		t.Fatalf("expected ErrNotSupported, got %v", err)
+	}
+}

github/rate_limit.go

@@ -0,0 +1,25 @@
+package github
+
+import (
+	"context"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+func (f *gitHubForge) GetRateLimit(ctx context.Context) (*forge.RateLimit, error) {
+	limits, _, err := f.client.RateLimit.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if limits == nil || limits.Core == nil {
+		return &forge.RateLimit{}, nil
+	}
+
+	core := limits.Core
+	return &forge.RateLimit{
+		Limit:     core.Limit,
+		Remaining: core.Remaining,
+		Reset:     core.Reset.Time,
+	}, nil
+}

github/rate_limit_test.go

@@ -0,0 +1,47 @@
+package github
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/google/go-github/v82/github"
+)
+
+func TestGitHubGetRateLimit(t *testing.T) {
+	resetTime := time.Date(2024, 6, 1, 12, 0, 0, 0, time.UTC)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("GET /api/v3/rate_limit", func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"resources": map[string]any{
+				"core": map[string]any{
+					"limit":     5000,
+					"remaining": 4999,
+					"reset":     resetTime.Unix(),
+				},
+			},
+		})
+	})
+
+	srv := httptest.NewServer(mux)
+	defer srv.Close()
+
+	c := github.NewClient(nil)
+	c, _ = c.WithEnterpriseURLs(srv.URL+"/api/v3", srv.URL+"/api/v3")
+	f := &gitHubForge{client: c}
+
+	rl, err := f.GetRateLimit(context.Background())
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	assertEqualInt(t, "Limit", 5000, rl.Limit)
+	assertEqualInt(t, "Remaining", 4999, rl.Remaining)
+	if !rl.Reset.Equal(resetTime) {
+		t.Errorf("Reset: want %v, got %v", resetTime, rl.Reset)
+	}
+}

gitlab/rate_limit.go

@@ -0,0 +1,33 @@
+package gitlab
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	forge "github.com/git-pkgs/forge"
+)
+
+func (f *gitLabForge) GetRateLimit(ctx context.Context) (*forge.RateLimit, error) {
+	// GitLab has no dedicated rate limit endpoint. Rate limit info comes
+	// from response headers on any API call, so we make a lightweight request.
+	_, resp, err := f.client.Version.GetVersion()
+	if err != nil {
+		return nil, err
+	}
+
+	limit, _ := strconv.Atoi(resp.Header.Get("RateLimit-Limit"))
+	remaining, _ := strconv.Atoi(resp.Header.Get("RateLimit-Remaining"))
+	resetUnix, _ := strconv.ParseInt(resp.Header.Get("RateLimit-Reset"), 10, 64)
+
+	var reset time.Time
+	if resetUnix > 0 {
+		reset = time.Unix(resetUnix, 0)
+	}
+
+	return &forge.RateLimit{
+		Limit:     limit,
+		Remaining: remaining,
+		Reset:     reset,
+	}, nil
+}