feat: Allow plugins to be downloaded using a user team_name context (…

…#162) Adding the ability to pass the `team_name` to the download. This will be passed in by the CLI as either the `team_name` in the user's saved configuration (i.e. `cloudquery switch` config) or the `team_name` associated with the API key being used. Ref: cloudquery/cloudquery-issues#835
cloudquery · Nov 14, 2023 · 5e33845 · 5e33845
1 parent 3c944dd
commit 5e33845
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 32 deletions.
diff --git a/managedplugin/download.go b/managedplugin/download.go
@@ -92,62 +92,59 @@ func getURLLocation(ctx context.Context, org string, name string, version string
 	return "", fmt.Errorf("failed to find plugin %s/%s version %s", org, name, version)
 }
 
-func DownloadPluginFromHub(ctx context.Context, authToken, localPath, team, name, version string, typ PluginType) error {
-	downloadDir := filepath.Dir(localPath)
-	if _, err := os.Stat(localPath); err == nil {
+type HubDownloadOptions struct {
+	AuthToken     string
+	TeamName      string
+	LocalPath     string
+	PluginTeam    string
+	PluginKind    string
+	PluginName    string
+	PluginVersion string
+}
+
+func DownloadPluginFromHub(ctx context.Context, ops HubDownloadOptions) error {
+	downloadDir := filepath.Dir(ops.LocalPath)
+	if _, err := os.Stat(ops.LocalPath); err == nil {
 		return nil
 	}
 
 	if err := os.MkdirAll(downloadDir, 0755); err != nil {
 		return fmt.Errorf("failed to create plugin directory %s: %w", downloadDir, err)
 	}
 
-	target := fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH)
-	c, err := cloudquery_api.NewClientWithResponses(APIBaseURL(),
-		cloudquery_api.WithRequestEditorFn(func(ctx context.Context, req *http.Request) error {
-			if authToken != "" {
-				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", authToken))
-			}
-			return nil
-		}))
-	if err != nil {
-		return fmt.Errorf("failed to create Hub API client: %w", err)
-	}
-
-	aj := "application/json"
-	resp, err := c.DownloadPluginAssetWithResponse(ctx, team, cloudquery_api.PluginKind(typ.String()), name, version, target, &cloudquery_api.DownloadPluginAssetParams{Accept: &aj})
+	pluginAsset, statusCode, err := downloadPluginAssetFromHub(ctx, ops)
 	if err != nil {
 		return fmt.Errorf("failed to get plugin url: %w", err)
 	}
-	switch resp.StatusCode() {
+	switch statusCode {
 	case http.StatusOK:
 		// we allow this status code
 	case http.StatusUnauthorized:
 		return fmt.Errorf("unauthorized. Try logging in via `cloudquery login`")
 	case http.StatusNotFound:
-		return fmt.Errorf("failed to download plugin %v %v/%v@%v: plugin version not found. If you're trying to use a private plugin you'll need to run `cloudquery login` first", typ, team, name, version)
+		return fmt.Errorf("failed to download plugin %v %v/%v@%v: plugin version not found. If you're trying to use a private plugin you'll need to run `cloudquery login` first", ops.PluginKind, ops.PluginTeam, ops.PluginName, ops.PluginVersion)
 	case http.StatusTooManyRequests:
 		return fmt.Errorf("too many download requests. Try logging in via `cloudquery login` to increase rate limits")
 	default:
-		return fmt.Errorf("failed to download plugin %v %v/%v@%v: unexpected status code %v", typ, team, name, version, resp.StatusCode())
+		return fmt.Errorf("failed to download plugin %v %v/%v@%v: unexpected status code %v", ops.PluginKind, ops.PluginTeam, ops.PluginName, ops.PluginVersion, statusCode)
 	}
-	if resp.JSON200 == nil {
-		return fmt.Errorf("failed to get plugin url for %v %v/%v@%v: missing json response", typ, team, name, version)
+	if pluginAsset == nil {
+		return fmt.Errorf("failed to get plugin url for %v %v/%v@%v: missing json response", ops.PluginKind, ops.PluginTeam, ops.PluginName, ops.PluginVersion)
 	}
-	location := resp.JSON200.Location
+	location := pluginAsset.Location
 	if len(location) == 0 {
 		return fmt.Errorf("failed to get plugin url: empty location from response")
 	}
-	pluginZipPath := localPath + ".zip"
+	pluginZipPath := ops.LocalPath + ".zip"
 	writtenChecksum, err := downloadFile(ctx, pluginZipPath, location)
 	if err != nil {
 		return fmt.Errorf("failed to download plugin: %w", err)
 	}
 
-	if resp.JSON200.Checksum == "" {
+	if pluginAsset.Checksum == "" {
 		fmt.Printf("Warning - checksum not verified: %s\n", writtenChecksum)
-	} else if writtenChecksum != resp.JSON200.Checksum {
-		return fmt.Errorf("checksum mismatch: expected %s, got %s", resp.JSON200.Checksum, writtenChecksum)
+	} else if writtenChecksum != pluginAsset.Checksum {
+		return fmt.Errorf("checksum mismatch: expected %s, got %s", pluginAsset.Checksum, writtenChecksum)
 	}
 
 	archive, err := zip.OpenReader(pluginZipPath)
@@ -156,14 +153,14 @@ func DownloadPluginFromHub(ctx context.Context, authToken, localPath, team, name
 	}
 	defer archive.Close()
 
-	fileInArchive, err := archive.Open(fmt.Sprintf("plugin-%s-%s-%s-%s", name, version, runtime.GOOS, runtime.GOARCH))
+	fileInArchive, err := archive.Open(fmt.Sprintf("plugin-%s-%s-%s-%s", ops.PluginName, ops.PluginVersion, runtime.GOOS, runtime.GOARCH))
 	if err != nil {
 		return fmt.Errorf("failed to open plugin archive: %w", err)
 	}
 
-	out, err := os.OpenFile(localPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0744)
+	out, err := os.OpenFile(ops.LocalPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0744)
 	if err != nil {
-		return fmt.Errorf("failed to create file %s: %w", localPath, err)
+		return fmt.Errorf("failed to create file %s: %w", ops.LocalPath, err)
 	}
 	_, err = io.Copy(out, fileInArchive)
 	if err != nil {
@@ -176,6 +173,54 @@ func DownloadPluginFromHub(ctx context.Context, authToken, localPath, team, name
 	return nil
 }
 
+func downloadPluginAssetFromHub(ctx context.Context, ops HubDownloadOptions) (*cloudquery_api.PluginAsset, int, error) {
+	c, err := cloudquery_api.NewClientWithResponses(APIBaseURL(),
+		cloudquery_api.WithRequestEditorFn(func(ctx context.Context, req *http.Request) error {
+			if ops.AuthToken != "" {
+				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", ops.AuthToken))
+			}
+			return nil
+		}))
+	if err != nil {
+		return nil, -1, fmt.Errorf("failed to create Hub API client: %w", err)
+	}
+
+	target := fmt.Sprintf("%s_%s", runtime.GOOS, runtime.GOARCH)
+	aj := "application/json"
+
+	switch {
+	case ops.TeamName != "":
+		resp, err := c.DownloadPluginAssetByTeamWithResponse(
+			ctx,
+			ops.TeamName,
+			ops.PluginTeam,
+			cloudquery_api.PluginKind(ops.PluginKind),
+			ops.PluginName,
+			ops.PluginVersion,
+			target,
+			&cloudquery_api.DownloadPluginAssetByTeamParams{Accept: &aj},
+		)
+		if err != nil {
+			return nil, -1, fmt.Errorf("failed to get plugin url with team: %w", err)
+		}
+		return resp.JSON200, resp.StatusCode(), nil
+	default:
+		resp, err := c.DownloadPluginAssetWithResponse(
+			ctx,
+			ops.PluginTeam,
+			cloudquery_api.PluginKind(ops.PluginKind),
+			ops.PluginName,
+			ops.PluginVersion,
+			target,
+			&cloudquery_api.DownloadPluginAssetParams{Accept: &aj},
+		)
+		if err != nil {
+			return nil, -1, fmt.Errorf("failed to get plugin url: %w", err)
+		}
+		return resp.JSON200, resp.StatusCode(), nil
+	}
+}
+
 func DownloadPluginFromGithub(ctx context.Context, localPath string, org string, name string, version string, typ PluginType) error {
 	downloadDir := filepath.Dir(localPath)
 	pluginZipPath := localPath + ".zip"

diff --git a/managedplugin/download_test.go b/managedplugin/download_test.go
@@ -48,7 +48,15 @@ func TestDownloadPluginFromCloudQueryHub(t *testing.T) {
 	}
 	for _, tc := range cases {
 		t.Run(tc.testName, func(t *testing.T) {
-			err := DownloadPluginFromHub(context.Background(), "", path.Join(tmp, tc.testName), tc.team, tc.plugin, tc.version, tc.typ)
+			err := DownloadPluginFromHub(context.Background(), HubDownloadOptions{
+				LocalPath:     path.Join(tmp, tc.testName),
+				AuthToken:     "",
+				TeamName:      "",
+				PluginTeam:    tc.team,
+				PluginKind:    tc.typ.String(),
+				PluginName:    tc.plugin,
+				PluginVersion: tc.version,
+			})
 			if (err != nil) != tc.wantErr {
 				t.Errorf("TestDownloadPluginFromCloudQueryIntegration() error = %v, wantErr %v", err, tc.wantErr)
 				return

diff --git a/managedplugin/options.go b/managedplugin/options.go
@@ -45,3 +45,9 @@ func WithAuthToken(authToken string) Option {
 		c.authToken = authToken
 	}
 }
+
+func WithTeamName(teamName string) Option {
+	return func(c *Client) {
+		c.teamName = teamName
+	}
+}
diff --git a/managedplugin/plugin.go b/managedplugin/plugin.go
@@ -91,6 +91,7 @@ type Client struct {
 	metrics              *Metrics
 	registry             Registry
 	authToken            string
+	teamName             string
 }
 
 // typ will be deprecated soon but now required for a transition period
@@ -182,7 +183,15 @@ func (c *Client) downloadPlugin(ctx context.Context, typ PluginType) error {
 		org, name := pathSplit[0], pathSplit[1]
 		c.LocalPath = filepath.Join(c.directory, "plugins", typ.String(), org, name, c.config.Version, "plugin")
 		c.LocalPath = WithBinarySuffix(c.LocalPath)
-		return DownloadPluginFromHub(ctx, c.authToken, c.LocalPath, org, name, c.config.Version, typ)
+		return DownloadPluginFromHub(ctx, HubDownloadOptions{
+			AuthToken:     c.authToken,
+			TeamName:      c.teamName,
+			LocalPath:     c.LocalPath,
+			PluginTeam:    org,
+			PluginKind:    typ.String(),
+			PluginName:    name,
+			PluginVersion: c.config.Version,
+		})
 	default:
 		return fmt.Errorf("unknown registry %s", c.config.Registry.String())
 	}