Merge pull request #332 from depot/feat/load-from-registry

feat: allow loading build from the depot registry

Author: ischolten

Diff for: pkg/build/build.go

@@ -35,6 +35,12 @@ type Credential struct {
 	Token string
 }
 
+type PullInfo struct {
+	Reference string
+	Username  string
+	Password  string
+}
+
 func (b *Build) AdditionalTags() []string {
 	if b.Response == nil || b.Response.Msg == nil {
 		return []string{fmt.Sprintf("registry.depot.dev/%s:%s", b.projectID, b.ID)}
@@ -92,6 +98,19 @@ func (b *Build) BuildProject() string {
 	return b.Response.Msg.ProjectId
 }
 
+func (b *Build) LoadUsingRegistry() bool {
+	if b.Response == nil || b.Response.Msg == nil {
+		return false
+	}
+
+	registry := b.Response.Msg.GetRegistry()
+	if registry == nil {
+		return false
+	}
+
+	return registry.LoadUsingRegistry
+}
+
 func NewBuild(ctx context.Context, req *cliv1.CreateBuildRequest, token string) (Build, error) {
 	client := depotapi.NewBuildClient()
 	res, err := client.CreateBuild(ctx, depotapi.WithAuthentication(connect.NewRequest(req), token))
@@ -157,6 +176,22 @@ func FromExistingBuild(ctx context.Context, buildID, token string, buildRes *con
 
 }
 
+func PullBuildInfo(ctx context.Context, buildID, token string) (*PullInfo, error) {
+	// Download location and credentials of image save.
+	client := depotapi.NewBuildClient()
+	req := &cliv1.GetPullInfoRequest{BuildId: buildID}
+	res, err := client.GetPullInfo(ctx, depotapi.WithAuthentication(connect.NewRequest(req), token))
+	if err != nil {
+		return nil, err
+	}
+
+	return &PullInfo{
+		Reference: res.Msg.Reference,
+		Username:  res.Msg.Username,
+		Password:  res.Msg.Password,
+	}, nil
+}
+
 type authProvider struct {
 	credentials map[string]clitypes.AuthConfig
 	dockerAuth  driver.Auth

Diff for: pkg/buildx/commands/bake.go

@@ -10,6 +10,7 @@ import (
 	"sync"
 
 	"github.com/containerd/containerd/platforms"
+	depotbuild "github.com/depot/cli/pkg/build"
 	"github.com/depot/cli/pkg/buildx/bake"
 	"github.com/depot/cli/pkg/buildx/build"
 	"github.com/depot/cli/pkg/buildx/builder"
@@ -102,6 +103,8 @@ func RunBake(dockerCli command.Cli, in BakeOptions, validator BakeValidator, pri
 				BuildID:      in.DepotOptions.buildID,
 				IsBake:       true,
 				ProgressMode: in.progress,
+				UseRegistry:  in.DepotOptions.loadUsingRegistry,
+				PullInfo:     in.DepotOptions.pullInfo,
 			},
 		)
 	}
@@ -179,7 +182,11 @@ func RunBake(dockerCli command.Cli, in BakeOptions, validator BakeValidator, pri
 					if slices.Contains(requestedTargets, resp[i].Name) {
 						reportingPrinter := progresshelper.NewReporter(ctx2, printer, in.buildID, in.token)
 						defer reportingPrinter.Close()
-						err = load.DepotFastLoad(ctx2, dockerCli.Client(), depotResponses, pullOpts, reportingPrinter)
+						if in.DepotOptions.loadUsingRegistry && in.DepotOptions.pullInfo != nil {
+							err = load.DepotLoadFromRegistry(ctx, dockerCli.Client(), in.DepotOptions.pullInfo.Reference, true, pullOpts, reportingPrinter)
+						} else {
+							err = load.DepotFastLoad(ctx2, dockerCli.Client(), depotResponses, pullOpts, reportingPrinter)
+						}
 					}
 					load.DeleteExportLeases(ctx2, depotResponses)
 					return err
@@ -318,6 +325,16 @@ func BakeCmd() *cobra.Command {
 				if buildProject != "" {
 					options.project = buildProject
 				}
+				loadUsingRegistry := build.LoadUsingRegistry()
+				if options.exportLoad && loadUsingRegistry {
+					options.save = true
+					pullInfo, err := depotbuild.PullBuildInfo(context.Background(), build.ID, token)
+					// if we cannot get pull info, dont fail; load as normal
+					if err == nil {
+						options.loadUsingRegistry = loadUsingRegistry
+						options.pullInfo = pullInfo
+					}
+				}
 				if options.save {
 					options.additionalCredentials = build.AdditionalCredentials()
 					options.additionalTags = build.AdditionalTags()

Diff for: pkg/buildx/commands/build.go

@@ -123,6 +123,8 @@ type DepotOptions struct {
 	saveTags              []string
 	additionalTags        []string
 	additionalCredentials []depotbuild.Credential
+	loadUsingRegistry     bool
+	pullInfo              *depotbuild.PullInfo
 
 	lint       bool
 	lintFailOn string
@@ -232,6 +234,8 @@ func buildTargets(ctx context.Context, dockerCli command.Cli, nodes []builder.No
 				BuildID:      depotOpts.buildID,
 				IsBake:       false,
 				ProgressMode: progressMode,
+				UseRegistry:  depotOpts.loadUsingRegistry,
+				PullInfo:     depotOpts.pullInfo,
 			},
 		)
 	}
@@ -324,7 +328,13 @@ func buildTargets(ctx context.Context, dockerCli command.Cli, nodes []builder.No
 
 	// NOTE: the err is returned at the end of this function after the final prints.
 	reportingPrinter := progresshelper.NewReporter(ctx, printer, depotOpts.buildID, depotOpts.token)
-	err = load.DepotFastLoad(ctx, dockerCli.Client(), resp, pullOpts, reportingPrinter)
+
+	if depotOpts.loadUsingRegistry && depotOpts.pullInfo != nil {
+		err = load.DepotLoadFromRegistry(ctx, dockerCli.Client(), depotOpts.pullInfo.Reference, false, pullOpts, reportingPrinter)
+	} else {
+		err = load.DepotFastLoad(ctx, dockerCli.Client(), resp, pullOpts, reportingPrinter)
+	}
+
 	if err != nil && !errors.Is(err, context.Canceled) {
 		// For now, we will fallback by rebuilding with load.
 		if exportLoad {
@@ -680,6 +690,16 @@ func BuildCmd() *cobra.Command {
 			if buildProject != "" {
 				options.project = buildProject
 			}
+			loadUsingRegistry := build.LoadUsingRegistry()
+			if options.exportLoad && loadUsingRegistry {
+				options.save = true
+				pullInfo, err := depotbuild.PullBuildInfo(context.Background(), build.ID, token)
+				// if we cannot get pull info, dont fail; load as normal
+				if err == nil {
+					options.loadUsingRegistry = loadUsingRegistry
+					options.pullInfo = pullInfo
+				}
+			}
 			if options.save {
 				options.additionalCredentials = build.AdditionalCredentials()
 				options.additionalTags = build.AdditionalTags()

Diff for: pkg/cmd/pull/options.go

@@ -12,13 +12,13 @@ import (
 	"golang.org/x/exp/slices"
 )
 
-func isSavedBuild(options []*cliv1.BuildOptions) bool {
+func isSavedBuild(options []*cliv1.BuildOptions, savedForLoad bool) bool {
 	for _, opt := range options {
 		if opt.Save {
 			return true
 		}
 	}
-	return false
+	return savedForLoad
 }
 
 func isBake(options []*cliv1.BuildOptions) bool {

Diff for: pkg/cmd/pull/pull.go

@@ -103,7 +103,8 @@ func NewCmdPull() *cobra.Command {
 			}
 
 			buildOptions := res.Msg.Options
-			if len(buildOptions) > 0 && !isSavedBuild(buildOptions) {
+			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)
 			}
 

Diff for: pkg/load/cli.go

@@ -4,17 +4,20 @@ import (
 	"fmt"
 	"strings"
 
+	depotbuild "github.com/depot/cli/pkg/build"
 	"github.com/docker/buildx/build"
 	"github.com/docker/buildx/util/progress"
 	"github.com/moby/buildkit/client"
 )
 
 // DepotLoadOptions are options to load images from the depot hosted registry.
 type DepotLoadOptions struct {
-	Project      string // Depot project name; used to tag images.
-	BuildID      string // Depot build ID; used to tag images.
-	IsBake       bool   // If run from bake, we add the bake target to the image tag.
-	ProgressMode string // ProgressMode quiet will not print progress.
+	Project      string               // Depot project name; used to tag images.
+	BuildID      string               // Depot build ID; used to tag images.
+	IsBake       bool                 // If run from bake, we add the bake target to the image tag.
+	ProgressMode string               // ProgressMode quiet will not print progress.
+	UseRegistry  bool                 // If UseRegistry, load build from registry instead of proxy
+	PullInfo     *depotbuild.PullInfo // If UseRegistry, the credentials for pulling from registry
 }
 
 // Options to download from the Depot hosted registry and tag the image with the user provide tag.
@@ -69,6 +72,15 @@ func WithDepotImagePull(buildOpts map[string]build.Options, loadOpts DepotLoadOp
 				UserTags: userTags,
 				Quiet:    loadOpts.ProgressMode == progress.PrinterModeQuiet,
 			}
+
+			if loadOpts.UseRegistry && loadOpts.PullInfo != nil {
+				serverAddress := "registry.depot.dev"
+				pullOpt.KeepImage = true
+				pullOpt.Username = &loadOpts.PullInfo.Username
+				pullOpt.Password = &loadOpts.PullInfo.Password
+				pullOpt.ServerAddress = &serverAddress
+			}
+
 			toPull[target] = pullOpt
 		}
 	}

Diff for: pkg/load/load.go

@@ -75,6 +75,29 @@ func DepotFastLoad(ctx context.Context, dockerapi docker.APIClient, resp []depot
 	return nil
 }
 
+func DepotLoadFromRegistry(ctx context.Context, dockerapi docker.APIClient, reference string, isBake bool, pullOpts map[string]PullOptions, printer progress.Writer) error {
+	if len(pullOpts) == 0 {
+		return nil
+	}
+
+	for target, pullOpt := range pullOpts {
+		pw := progress.WithPrefix(printer, target, len(pullOpts) > 1)
+
+		imageName := reference
+
+		if isBake {
+			imageName = fmt.Sprintf("%s-%s", reference, target)
+		}
+
+		err := PullImages(ctx, dockerapi, imageName, pullOpt, pw)
+		if err != nil {
+			return fmt.Errorf("failed to pull image: %w", err)
+		}
+	}
+
+	return nil
+}
+
 // For now if there is a multi-platform build we try to only download the
 // architecture of the depot CLI host.  If there is not a node with the same
 // architecture as the  depot CLI host, we take the first node in the list.