add rwx artifacts download command (#274)

Author: tonywok

cmd/rwx/artifacts.go

@@ -0,0 +1,18 @@
+package main
+
+import (
+	"github.com/rwx-cloud/cli/cmd/rwx/artifacts"
+	"github.com/rwx-cloud/cli/internal/cli"
+	"github.com/spf13/cobra"
+)
+
+var artifactsCmd = &cobra.Command{
+	GroupID: "results",
+	Use:     "artifacts",
+	Short:   "Manage task artifacts",
+}
+
+func init() {
+	artifacts.InitDownload(requireAccessToken, func() cli.Service { return service })
+	artifactsCmd.AddCommand(artifacts.DownloadCmd)
+}

cmd/rwx/artifacts/download.go

@@ -0,0 +1,104 @@
+package artifacts
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+
+	"github.com/rwx-cloud/cli/internal/cli"
+	"github.com/rwx-cloud/cli/internal/errors"
+
+	"github.com/spf13/cobra"
+)
+
+var (
+	downloadOutputDir   string
+	downloadOutputFile  string
+	downloadJSON        bool
+	downloadAutoExtract bool
+	downloadOpen        bool
+
+	DownloadCmd *cobra.Command
+)
+
+func InitDownload(requireAccessToken func() error, getService func() cli.Service) {
+	DownloadCmd = &cobra.Command{
+		Args: cobra.ExactArgs(2),
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return requireAccessToken()
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			taskID := args[0]
+			artifactKey := args[1]
+
+			outputDirSet := cmd.Flags().Changed("output-dir")
+			outputFileSet := cmd.Flags().Changed("output-file")
+			if outputDirSet && outputFileSet {
+				return errors.New("output-dir and output-file cannot be used together")
+			}
+
+			var absOutputDir string
+			var absOutputFile string
+			var err error
+
+			if downloadOutputFile != "" {
+				absOutputFile, err = filepath.Abs(downloadOutputFile)
+				if err != nil {
+					return errors.Wrapf(err, "unable to resolve absolute path for %s", downloadOutputFile)
+				}
+			} else {
+				outputDir := downloadOutputDir
+				if !outputDirSet {
+					outputDir, err = getDefaultDownloadsDir()
+					if err != nil {
+						return errors.Wrap(err, "unable to determine default downloads directory")
+					}
+				}
+				absOutputDir, err = filepath.Abs(outputDir)
+				if err != nil {
+					return errors.Wrapf(err, "unable to resolve absolute path for %s", outputDir)
+				}
+			}
+
+			err = getService().DownloadArtifact(cli.DownloadArtifactConfig{
+				TaskID:                 taskID,
+				ArtifactKey:            artifactKey,
+				OutputDir:              absOutputDir,
+				OutputFile:             absOutputFile,
+				OutputDirExplicitlySet: outputDirSet,
+				Json:                   downloadJSON,
+				AutoExtract:            downloadAutoExtract,
+				Open:                   downloadOpen,
+			})
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+		Short: "Download an artifact from a task",
+		Use:   "download <task-id> <artifact-key> [flags]",
+	}
+
+	DownloadCmd.Flags().StringVar(&downloadOutputDir, "output-dir", "", "output directory for the downloaded artifact (defaults to Downloads folder)")
+	DownloadCmd.Flags().StringVar(&downloadOutputFile, "output-file", "", "output file path for the downloaded artifact")
+	DownloadCmd.MarkFlagsMutuallyExclusive("output-dir", "output-file")
+	DownloadCmd.Flags().BoolVar(&downloadJSON, "json", false, "output file locations as JSON")
+	DownloadCmd.Flags().BoolVar(&downloadAutoExtract, "auto-extract", false, "automatically extract directory tar archives")
+	DownloadCmd.Flags().BoolVar(&downloadOpen, "open", false, "automatically open the downloaded file(s)")
+}
+
+func getDefaultDownloadsDir() (string, error) {
+	if runtime.GOOS == "linux" {
+		if xdgDownload := os.Getenv("XDG_DOWNLOAD_DIR"); xdgDownload != "" {
+			return xdgDownload, nil
+		}
+	}
+
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return "", errors.Wrap(err, "unable to determine user home directory")
+	}
+
+	return filepath.Join(homeDir, "Downloads"), nil
+}

cmd/rwx/root.go

@@ -117,6 +117,7 @@ func init() {
 	rootCmd.SetCompletionCommandGroupID("setup")
 
 	// Add commands (GroupID is set in each command's definition)
+	rootCmd.AddCommand(artifactsCmd)
 	rootCmd.AddCommand(debugCmd)
 	rootCmd.AddCommand(dispatchCmd)
 	rootCmd.AddCommand(imageCmd)

internal/api/client.go

@@ -701,6 +701,60 @@ func (c Client) DownloadLogs(request LogDownloadRequestResult, maxRetryDurationS
 	}
 }
 
+func (c Client) GetArtifactDownloadRequest(taskId, artifactKey string) (ArtifactDownloadRequestResult, error) {
+	endpoint := fmt.Sprintf("/mint/api/tasks/%s/artifact_downloads/%s", url.PathEscape(taskId), url.PathEscape(artifactKey))
+	result := ArtifactDownloadRequestResult{}
+
+	req, err := http.NewRequest(http.MethodGet, endpoint, nil)
+	if err != nil {
+		return result, errors.Wrap(err, "unable to create new HTTP request")
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.RoundTrip(req)
+	if err != nil {
+		return result, errors.Wrap(err, "HTTP request failed")
+	}
+	defer resp.Body.Close()
+
+	if err = decodeResponseJSON(resp, &result); err != nil {
+		return result, err
+	}
+
+	return result, nil
+}
+
+func (c Client) DownloadArtifact(request ArtifactDownloadRequestResult) ([]byte, error) {
+	req, err := http.NewRequest(http.MethodGet, request.URL, nil)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to create new HTTP request")
+	}
+	req.Header.Set("Accept", "application/octet-stream")
+
+	// Use http.DefaultClient directly since the artifact will come from storage (S3, etc.)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, errors.Wrap(err, "HTTP request failed")
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
+		artifactBytes, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to read response body")
+		}
+		return artifactBytes, nil
+	}
+
+	bodyBytes, _ := io.ReadAll(resp.Body)
+	errMsg := extractErrorMessage(bytes.NewReader(bodyBytes))
+	if errMsg == "" {
+		errMsg = fmt.Sprintf("Unable to download artifact - %s", resp.Status)
+	}
+	return nil, errors.New(errMsg)
+}
+
 func decodeResponseJSON(resp *http.Response, result any) error {
 	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
 		if result == nil {