Merge pull request #391 from depot/feat/improve-depot-pull

feat: improve `depot pull` to support tags

Author: alexisbouchez

pkg/cmd/pull/pull.go

@@ -3,6 +3,7 @@ package pull
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"connectrpc.com/connect"
 	depotapi "github.com/depot/cli/pkg/api"
@@ -18,19 +19,23 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+const (
+	depotRegistry = "registry.depot.dev"
+)
+
 func NewCmdPull() *cobra.Command {
 	var (
 		token     string
 		projectID string
 		platform  string
-		buildID   string
+		param     string
 		progress  string
 		userTags  []string
 		targets   []string
 	)
 
 	cmd := &cobra.Command{
-		Use:   "pull [flags] [buildID]",
+		Use:   "pull [flags] [buildID|tag]",
 		Short: "Pull a project's build from the Depot registry",
 		Args:  cli.RequiresMaxArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -40,7 +45,7 @@ func NewCmdPull() *cobra.Command {
 			}
 
 			if len(args) > 0 {
-				buildID = args[0]
+				param = args[0]
 			}
 			_, isCI := ci.Provider()
 			if progress == prog.PrinterModeAuto && isCI {
@@ -49,7 +54,7 @@ func NewCmdPull() *cobra.Command {
 
 			ctx := cmd.Context()
 
-			token, err := helpers.ResolveProjectAuth(ctx, token)
+			token, err = helpers.ResolveProjectAuth(ctx, token)
 			if err != nil {
 				return err
 			}
@@ -58,7 +63,7 @@ func NewCmdPull() *cobra.Command {
 				return fmt.Errorf("missing API token, please run `depot login`")
 			}
 
-			if buildID == "" {
+			if param == "" {
 				var selectedProject *helpers.SelectedProject
 				projectID = helpers.ResolveProjectID(projectID)
 				if projectID == "" { // No locally saved depot.json.
@@ -85,34 +90,53 @@ func NewCmdPull() *cobra.Command {
 					return fmt.Errorf("build ID must be specified")
 				}
 
-				buildID, err = helpers.SelectBuildID(ctx, token, projectID, client)
+				param, err = helpers.SelectBuildID(ctx, token, projectID, client)
 				if err != nil {
 					return err
 				}
 
-				if buildID == "" {
-					return fmt.Errorf("build ID must be specified")
+				if param == "" {
+					return fmt.Errorf("build ID or tag must be specified")
 				}
 			}
 
+			// Check if the buildID is actually a registry reference or tag
+			if strings.HasPrefix(param, depotRegistry+"/") {
+				// Extract project ID and tag from the reference
+				projectID, tag := extractProjectIDAndTag(param)
+				return pullByTag(ctx, dockerCli, token, projectID, tag, userTags, platform, progress)
+			}
+
+			// Check if the param is in the format "projectID:tag"
+			if strings.Contains(param, ":") && !strings.HasPrefix(param, depotRegistry+"/") {
+				parts := strings.SplitN(param, ":", 2)
+				if len(parts) == 2 {
+					projectID = parts[0]
+					tag := parts[1]
+					return pullByTag(ctx, dockerCli, token, projectID, tag, userTags, platform, progress)
+				}
+			}
+
+			// Try to get build info first (build ID approach)
 			client := depotapi.NewBuildClient()
-			req := &cliv1.GetPullInfoRequest{BuildId: buildID}
+			req := &cliv1.GetPullInfoRequest{BuildId: param}
 			res, err := client.GetPullInfo(ctx, depotapi.WithAuthentication(connect.NewRequest(req), token))
 			if err != nil {
-				return err
+				// If GetPullInfo fails, try the direct tag approach
+				return pullByTag(ctx, dockerCli, token, projectID, param, userTags, platform, progress)
 			}
 
 			buildOptions := res.Msg.Options
 			savedForLoad := res.Msg.SaveForLoad
 			if len(buildOptions) > 0 && !isSavedBuild(buildOptions, savedForLoad) {
-				return fmt.Errorf("build %s is not a saved build. To use the registry use --save when building", buildID)
+				return fmt.Errorf("build %s is not a saved build. To use the registry use --save when building", param)
 			}
 
 			if isBake(buildOptions) {
 				return pullBake(ctx, dockerCli, res.Msg, targets, userTags, platform, progress)
-			} else {
-				return pullBuild(ctx, dockerCli, res.Msg, userTags, platform, progress)
 			}
+
+			return pullBuild(ctx, dockerCli, res.Msg, userTags, platform, progress)
 		},
 	}
 
@@ -126,6 +150,68 @@ func NewCmdPull() *cobra.Command {
 	return cmd
 }
 
+// extractProjectIDAndTag extracts project ID and tag from a registry reference
+func extractProjectIDAndTag(reference string) (projectID string, tag string) {
+	// Remove the registry prefix
+	projectAndTag := strings.TrimPrefix(reference, depotRegistry+"/")
+
+	// Split on colon to separate project ID and tag
+	parts := strings.SplitN(projectAndTag, ":", 2)
+	if len(parts) != 2 {
+		return "", reference
+	}
+
+	return parts[0], parts[1]
+}
+
+// pullByTag handles pulling an image directly by tag when build ID approach fails
+func pullByTag(ctx context.Context, dockerCli command.Cli, token, projectID, tag string, userTags []string, platform, progress string) error {
+	client := depotapi.NewBuildClient()
+
+	// Get pull token for the project
+	req := &cliv1.GetPullTokenRequest{ProjectId: &projectID}
+	res, err := client.GetPullToken(ctx, depotapi.WithAuthentication(connect.NewRequest(req), token))
+	if err != nil {
+		return fmt.Errorf("failed to get pull token for project %s: %w", projectID, err)
+	}
+
+	// Construct the full image reference
+	imageName := fmt.Sprintf("%s/%s:%s", depotRegistry, projectID, tag)
+
+	// Set up pull options
+	serverAddress := depotRegistry
+	username := "x-token"
+	opts := load.PullOptions{
+		UserTags:      userTags,
+		Quiet:         progress == prog.PrinterModeQuiet,
+		KeepImage:     true,
+		Username:      &username,
+		Password:      &res.Msg.Token,
+		ServerAddress: &serverAddress,
+	}
+	if platform != "" {
+		opts.Platform = &platform
+	}
+
+	// Create a simple pull struct
+	pull := &pull{
+		imageName:   imageName,
+		pullOptions: opts,
+	}
+
+	// Set up printer
+	printer, cancel, err := buildPrinter(ctx, pull, progress)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		cancel()
+		_ = printer.Wait()
+	}()
+
+	return load.PullImages(ctx, dockerCli.Client(), pull.imageName, pull.pullOptions, printer)
+}
+
 func pullBuild(ctx context.Context, dockerCli command.Cli, msg *cliv1.GetPullInfoResponse, userTags []string, platform string, progress string) error {
 	pull := buildPullOpt(msg, userTags, platform, progress)
 	printer, cancel, err := buildPrinter(ctx, pull, progress)