diff --git a/examples/simple_plugin/go.sum b/examples/simple_plugin/go.sum index c333bf9c77..629badd84f 100644 --- a/examples/simple_plugin/go.sum +++ b/examples/simple_plugin/go.sum @@ -88,6 +88,8 @@ github.com/cloudquery/arrow/go/v14 v14.0.0-20231029080147-50d3871d0804 h1:y4EwAG github.com/cloudquery/arrow/go/v14 v14.0.0-20231029080147-50d3871d0804/go.mod h1:TqWp9yvMb9yZSxFNiij6cmZefm+1jw3oZU0L0w9lT7E= github.com/cloudquery/cloudquery-api-go v1.6.0 h1:8yAbNW+njhGmJLEnh7c55WFgs4rnGsIh3koIMkrebxg= github.com/cloudquery/cloudquery-api-go v1.6.0/go.mod h1:03fojQg0UpdgqXZ9tzZ5gF5CPad/F0sok66bsX6u4RA= +github.com/cloudquery/cloudquery-api-go v1.5.1 h1:7CVbn8/A2bljPMjAUfFALI97ZnWOpRcAW/aCAr1FWYg= +github.com/cloudquery/cloudquery-api-go v1.5.1/go.mod h1:03fojQg0UpdgqXZ9tzZ5gF5CPad/F0sok66bsX6u4RA= github.com/cloudquery/plugin-pb-go v1.14.0 h1:q2eZzOLVUPBjDj6AG1Smc+xE3Usuc5JDaJxv6s4M60A= github.com/cloudquery/plugin-pb-go v1.14.0/go.mod h1:Ti4SRHVDau7oF1w0/Oimt1k45yJilD0oFE0Xr7Tf3AI= github.com/cloudquery/plugin-sdk/v2 v2.7.0 h1:hRXsdEiaOxJtsn/wZMFQC9/jPfU1MeMK3KF+gPGqm7U= diff --git a/premium/usage.go b/premium/usage.go index ef75af6482..d0ada1dda7 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -168,9 +168,10 @@ func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pl op(u) } + tokenClient := auth.NewTokenClient() + // Create a default api client if none was provided if u.apiClient == nil { - tokenClient := auth.NewTokenClient() ac, err := cqapi.NewClientWithResponses(u.url, cqapi.WithRequestEditorFn(func(ctx context.Context, req *http.Request) error { token, err := tokenClient.GetToken() if err != nil { @@ -187,14 +188,9 @@ func NewUsageClient(pluginTeam cqapi.PluginTeam, pluginKind cqapi.PluginKind, pl // Set team name from configuration if not provided if u.teamName == "" { - teamName, err := config.GetValue("team") - 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 == "" { - return nil, fmt.Errorf("team name not set. Hint: use `cloudquery switch `") + teamName, err := u.getTeamNameByTokenType(tokenClient.GetTokenType()) + if err != nil { + return nil, fmt.Errorf("failed to get team name: %w", err) } u.teamName = teamName } @@ -361,3 +357,33 @@ func (u *BatchUpdater) calculateRetryDuration(statusCode int, headers http.Heade func retryableStatusCode(statusCode int) bool { return statusCode == http.StatusTooManyRequests || statusCode == http.StatusServiceUnavailable } + +func (u *BatchUpdater) getTeamNameByTokenType(tokenType auth.TokenType) (string, error) { + switch tokenType { + case auth.BearerToken: + teamName, err := config.GetValue("team") + if errors.Is(err, os.ErrNotExist) { + return "", 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 "", fmt.Errorf("failed to get team name from config: %w", err) + } + if teamName == "" { + return "", fmt.Errorf("team name not set. Hint: use `cloudquery switch `") + } + return teamName, nil + case auth.APIKey: + resp, err := u.apiClient.ListTeamsWithResponse(context.Background(), &cqapi.ListTeamsParams{}) + if err != nil { + return "", fmt.Errorf("failed to list teams for API key: %w", err) + } + if resp.StatusCode() != http.StatusOK { + return "", fmt.Errorf("failed to list teams for API key, status code: %s", resp.Status()) + } + if len(resp.JSON200.Items) != 1 { + return "", fmt.Errorf("expected to find exactly one team for API key, found %d", len(resp.JSON200.Items)) + } + return resp.JSON200.Items[0].Name, nil + default: + return "", fmt.Errorf("unsupported token type: %v", tokenType) + } +}