6071dc74f1
This commit introduces a new bandwidth limiting feature to the `borg collect` command. The feature is implemented using a token bucket algorithm in a new `pkg/ratelimit` package. The rate limiter is integrated with the `http.Client` via a custom `http.RoundTripper`, and the feature is exposed to the user through a new `--bandwidth` flag on the `collect` command. The bandwidth limiting feature has been applied to the `website` and `github` collectors, and unit and integration tests have been added to verify the functionality. The following changes have been made: - Created a new `pkg/ratelimit` package with a token bucket implementation. - Integrated the rate limiter with `http.Client` using a custom `http.RoundTripper`. - Added a `--bandwidth` flag to the `collect` command. - Applied the bandwidth limit to the `website` and `github` collectors. - Added unit tests for the rate limiter and bandwidth parsing logic. - Added integration tests for the `collect website` and `collect github repo` commands. The following issues were encountered and were being addressed when the session ended: - Build errors in the `cmd` package, specifically in `cmd/all.go` and `cmd/all_test.go`. - The need for a `MockGithubClient` in the `mocks` package. - The `website` package needs to be refactored to reduce code duplication. - The rate limiter's performance can be improved. Co-authored-by: Snider <631881+Snider@users.noreply.github.com>
46 lines
1.2 KiB
Go
46 lines
1.2 KiB
Go
package cmd
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"github.com/Snider/Borg/pkg/github"
|
|
"github.com/Snider/Borg/pkg/ratelimit"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var (
|
|
// GithubClient is the github client used by the command. It can be replaced for testing.
|
|
GithubClient = func(client *http.Client) github.GithubClient {
|
|
return github.NewGithubClient(client)
|
|
}
|
|
)
|
|
|
|
var collectGithubReposCmd = &cobra.Command{
|
|
Use: "repos [user-or-org]",
|
|
Short: "Collects all public repositories for a user or organization",
|
|
Args: cobra.ExactArgs(1),
|
|
RunE: func(cmd *cobra.Command, args []string) error {
|
|
bandwidth, _ := cmd.Flags().GetString("bandwidth")
|
|
bytesPerSec, err := ratelimit.ParseBandwidth(bandwidth)
|
|
if err != nil {
|
|
return fmt.Errorf("invalid bandwidth: %w", err)
|
|
}
|
|
|
|
client := github.NewAuthenticatedClient(cmd.Context(), ratelimit.NewRateLimitedRoundTripper(http.DefaultTransport, bytesPerSec))
|
|
githubClient := GithubClient(client)
|
|
|
|
repos, err := githubClient.GetPublicRepos(cmd.Context(), args[0])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, repo := range repos {
|
|
fmt.Fprintln(cmd.OutOrStdout(), repo)
|
|
}
|
|
return nil
|
|
},
|
|
}
|
|
|
|
func init() {
|
|
collectGithubCmd.AddCommand(collectGithubReposCmd)
|
|
}
|