fix: Retrieve team for api key (#1372)

Retreive the team associated with an API key to allow usage tracking against a team name

fixes: cloudquery/cloudquery-issues#893

Author: mnorbury

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=

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 <team>`", 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 <team>`")
+		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 <team>`", 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 <team>`")
+		}
+		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)
+	}
+}