From 6c4942a22bfbb2ce3140e51f373d22de5f432332 Mon Sep 17 00:00:00 2001 From: Herman Schaaf Date: Wed, 1 Nov 2023 19:59:33 +0000 Subject: [PATCH 1/2] Infer team name if none specified --- premium/usage.go | 64 ++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 10 deletions(-) diff --git a/premium/usage.go b/premium/usage.go index 1e1c2becf4..8eb1005375 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -2,6 +2,7 @@ package premium import ( "context" + "errors" "fmt" "math/rand" "net/http" @@ -144,7 +145,7 @@ type BatchUpdater struct { isClosed bool } -func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pluginName cqapi.PluginName, ops ...UsageClientOptions) (*BatchUpdater, error) { +func NewUsageClient(ctx context.Context, pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pluginName cqapi.PluginName, ops ...UsageClientOptions) (*BatchUpdater, error) { u := &BatchUpdater{ logger: zerolog.Nop(), url: defaultAPIURL, @@ -166,15 +167,6 @@ func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pl op(u) } - // Set team name from configuration if not provided - if u.teamName == "" { - teamName, err := config.GetValue("team") - if err != nil { - return nil, fmt.Errorf("failed to get team name from config: %w", err) - } - u.teamName = teamName - } - // Create a default api client if none was provided if u.apiClient == nil { tokenClient := auth.NewTokenClient() @@ -192,11 +184,63 @@ func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pl u.apiClient = ac } + // Set team name from configuration if not provided + if u.teamName == "" { + teamName, err := config.GetValue("team") + if err != nil { + return nil, fmt.Errorf("failed to get team name from config: %w", err) + } + if teamName == "" { + teamName, err = u.inferTeamNameFromAPI(ctx) + if err != nil { + return nil, fmt.Errorf("failed to infer team name from API: %w", err) + } + } + u.teamName = teamName + } + u.backgroundUpdater() return u, nil } +// inferTeamNameFromAPI fetches the team name from the API using the provided API token. If the user is part +// of exactly one team, this team will be returned. Otherwise, an error is returned. Users in +// multiple teams must specify the team name using `cloudquery switch `. +// This method works for API keys as well, because API keys are by definition tied to a single team. +func (u *BatchUpdater) inferTeamNameFromAPI(ctx context.Context) (string, error) { + u.logger.Debug().Msg("") + perPage := cqapi.PerPage(10) + params := cqapi.ListTeamsParams{ + PerPage: &perPage, + } + var err error + for try := 0; try < 3; try++ { + teams, err := u.apiClient.ListTeamsWithResponse(ctx, ¶ms) + if err != nil { + return "", fmt.Errorf("while listing teams: %w", err) + } + if teams.StatusCode() != http.StatusOK { + return "", fmt.Errorf("got non-200 status code listing teams: %d %s", teams.StatusCode(), teams.Status()) + } + if teams.StatusCode() == http.StatusTooManyRequests || teams.StatusCode() >= 500 { + u.logger.Warn().Int("code", teams.StatusCode()).Msg("got retryable code while listing teams, retrying in 2 seconds") + time.Sleep(2 * time.Second) + err = errors.Join(fmt.Errorf("got retryable code while listing teams: %v", teams.StatusCode())) + continue + } + items := teams.JSON200.Items + if len(items) == 0 { + return "", fmt.Errorf("no teams found for user") + } + if len(items) > 1 { + return "", fmt.Errorf("more than one team found. Hint: Try specifying a team using `cloudquery switch `") + } + return items[0].Name, nil + } + return "", err +} + func (u *BatchUpdater) Increase(rows uint32) error { if rows <= 0 { return fmt.Errorf("rows must be greater than zero got %d", rows) From e61e5911c3883090d94615161ea402688617c585 Mon Sep 17 00:00:00 2001 From: Herman Schaaf Date: Thu, 2 Nov 2023 10:00:12 +0000 Subject: [PATCH 2/2] Update --- premium/usage.go | 49 ++++++------------------------------------------ 1 file changed, 6 insertions(+), 43 deletions(-) diff --git a/premium/usage.go b/premium/usage.go index 8eb1005375..ef75af6482 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -6,6 +6,7 @@ import ( "fmt" "math/rand" "net/http" + "os" "sync/atomic" "time" @@ -145,7 +146,7 @@ type BatchUpdater struct { isClosed bool } -func NewUsageClient(ctx context.Context, pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pluginName cqapi.PluginName, ops ...UsageClientOptions) (*BatchUpdater, error) { +func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pluginName cqapi.PluginName, ops ...UsageClientOptions) (*BatchUpdater, error) { u := &BatchUpdater{ logger: zerolog.Nop(), url: defaultAPIURL, @@ -187,14 +188,13 @@ func NewUsageClient(ctx context.Context, pluginTeam cqapi.PluginTeam, pluginKind // Set team name from configuration if not provided if u.teamName == "" { teamName, err := config.GetValue("team") - if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, fmt.Errorf("config file for reading team name not found (%w). Hint: use `cloudquery login` and/or `cloudquery switch `", err) + } else if err != nil { return nil, fmt.Errorf("failed to get team name from config: %w", err) } if teamName == "" { - teamName, err = u.inferTeamNameFromAPI(ctx) - if err != nil { - return nil, fmt.Errorf("failed to infer team name from API: %w", err) - } + return nil, fmt.Errorf("team name not set. Hint: use `cloudquery switch `") } u.teamName = teamName } @@ -204,43 +204,6 @@ func NewUsageClient(ctx context.Context, pluginTeam cqapi.PluginTeam, pluginKind return u, nil } -// inferTeamNameFromAPI fetches the team name from the API using the provided API token. If the user is part -// of exactly one team, this team will be returned. Otherwise, an error is returned. Users in -// multiple teams must specify the team name using `cloudquery switch `. -// This method works for API keys as well, because API keys are by definition tied to a single team. -func (u *BatchUpdater) inferTeamNameFromAPI(ctx context.Context) (string, error) { - u.logger.Debug().Msg("") - perPage := cqapi.PerPage(10) - params := cqapi.ListTeamsParams{ - PerPage: &perPage, - } - var err error - for try := 0; try < 3; try++ { - teams, err := u.apiClient.ListTeamsWithResponse(ctx, ¶ms) - if err != nil { - return "", fmt.Errorf("while listing teams: %w", err) - } - if teams.StatusCode() != http.StatusOK { - return "", fmt.Errorf("got non-200 status code listing teams: %d %s", teams.StatusCode(), teams.Status()) - } - if teams.StatusCode() == http.StatusTooManyRequests || teams.StatusCode() >= 500 { - u.logger.Warn().Int("code", teams.StatusCode()).Msg("got retryable code while listing teams, retrying in 2 seconds") - time.Sleep(2 * time.Second) - err = errors.Join(fmt.Errorf("got retryable code while listing teams: %v", teams.StatusCode())) - continue - } - items := teams.JSON200.Items - if len(items) == 0 { - return "", fmt.Errorf("no teams found for user") - } - if len(items) > 1 { - return "", fmt.Errorf("more than one team found. Hint: Try specifying a team using `cloudquery switch `") - } - return items[0].Name, nil - } - return "", err -} - func (u *BatchUpdater) Increase(rows uint32) error { if rows <= 0 { return fmt.Errorf("rows must be greater than zero got %d", rows)