Add export and load for dagger images

Signed-off-by: Helder Correia <174525+helderco@users.noreply.github.com>

plan/task/export.go

@@ -0,0 +1,136 @@
+package task
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/moby/buildkit/exporter/containerimage/exptypes"
+
+	"github.com/docker/distribution/reference"
+	bk "github.com/moby/buildkit/client"
+	"github.com/moby/buildkit/client/llb"
+	"github.com/rs/zerolog/log"
+	"go.dagger.io/dagger/compiler"
+	"go.dagger.io/dagger/plancontext"
+	"go.dagger.io/dagger/solver"
+)
+
+func init() {
+	Register("Export", func() Task { return &exportTask{} })
+}
+
+type exportTask struct {
+}
+
+func (t exportTask) PreRun(ctx context.Context, pctx *plancontext.Context, v *compiler.Value) error {
+	dir, err := os.MkdirTemp("", "dagger-export-*")
+	if err != nil {
+		return err
+	}
+
+	pctx.TempDirs.Add(dir, v.Path().String())
+	pctx.LocalDirs.Add(dir)
+
+	return nil
+}
+
+func (t exportTask) Run(ctx context.Context, pctx *plancontext.Context, s solver.Solver, v *compiler.Value) (*compiler.Value, error) {
+	lg := log.Ctx(ctx)
+
+	dir := pctx.TempDirs.Get(v.Path().String())
+
+	var opts struct {
+		Tag  string
+		Path string
+		Type string
+	}
+
+	if err := v.Decode(&opts); err != nil {
+		return nil, err
+	}
+
+	switch opts.Type {
+	case bk.ExporterDocker, bk.ExporterOCI:
+	default:
+		return nil, fmt.Errorf("unsupported export type %q", opts.Type)
+	}
+
+	// Normalize tag
+	tag, err := reference.ParseNormalizedNamed(opts.Tag)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse ref %s: %w", opts.Tag, err)
+	}
+	tag = reference.TagNameOnly(tag)
+
+	lg.Debug().Str("tag", tag.String()).Msg("normalized tag")
+
+	// Get input state
+	input, err := pctx.FS.FromValue(v.Lookup("input"))
+	if err != nil {
+		return nil, err
+	}
+	st, err := input.State()
+	if err != nil {
+		return nil, err
+	}
+
+	// Decode the image config
+	imageConfig := ImageConfig{}
+	if err := v.Lookup("config").Decode(&imageConfig); err != nil {
+		return nil, err
+	}
+
+	img := NewImage(imageConfig, pctx.Platform.Get())
+
+	// Export image
+	resp, err := s.Export(ctx, st, &img, bk.ExportEntry{
+		Type: opts.Type,
+		Attrs: map[string]string{
+			"name": tag.String(),
+		},
+		Output: func(a map[string]string) (io.WriteCloser, error) {
+			file := filepath.Join(dir, opts.Path)
+			return os.Create(file)
+		},
+	}, pctx.Platform.Get())
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Save the image id
+	imageID, ok := resp.ExporterResponse[exptypes.ExporterImageConfigDigestKey]
+	if !ok {
+		return nil, fmt.Errorf("image export for %q did not return an image id", tag.String())
+	}
+
+	// FIXME: Remove the `Copy` and use `Local` directly.
+	//
+	// Copy'ing is a costly operation which should be unnecessary.
+	// However, using llb.Local directly breaks caching sometimes for unknown reasons.
+	outputState := llb.Scratch().File(
+		llb.Copy(
+			llb.Local(
+				dir,
+				withCustomName(v, "Export %s", opts.Path),
+			),
+			"/",
+			"/",
+		),
+		withCustomName(v, "Local %s [copy]", opts.Path),
+	)
+
+	result, err := s.Solve(ctx, outputState, pctx.Platform.Get())
+	if err != nil {
+		return nil, err
+	}
+
+	fs := pctx.FS.New(result)
+	return compiler.NewValue().FillFields(map[string]interface{}{
+		"output":  fs.MarshalCUE(),
+		"imageID": imageID,
+	})
+}

plan/task/image.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/moby/buildkit/frontend/dockerfile/dockerfile2llb"
 	"github.com/moby/buildkit/frontend/dockerfile/shell"
+	specs "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
 // ImageConfig defines the execution parameters which should be used as a base when running a container using an image.
@@ -152,3 +153,14 @@ func ConvertHealthConfig(spec *dockerfile2llb.HealthConfig) *HealthConfig {
 
 	return &cfg
 }
+
+func NewImage(config ImageConfig, platform specs.Platform) dockerfile2llb.Image {
+	return dockerfile2llb.Image{
+		Config: config.ToSpec(),
+		Image: specs.Image{
+			Architecture: platform.Architecture,
+			OS:           platform.OS,
+		},
+		Variant: platform.Variant,
+	}
+}

plan/task/push.go

@@ -6,7 +6,6 @@ import (
 
 	"github.com/docker/distribution/reference"
 	bk "github.com/moby/buildkit/client"
-	"github.com/moby/buildkit/frontend/dockerfile/dockerfile2llb"
 	"github.com/rs/zerolog/log"
 	"go.dagger.io/dagger/compiler"
 	"go.dagger.io/dagger/plancontext"
@@ -67,15 +66,11 @@ func (c *pushTask) Run(ctx context.Context, pctx *plancontext.Context, s solver.
 		return nil, err
 	}
 
-	// Add platform to image configuration
-	exportImageConfig := &dockerfile2llb.Image{Config: imageConfig.ToSpec()}
-	exportImageConfig.OS = pctx.Platform.Get().OS
-	exportImageConfig.Architecture = pctx.Platform.Get().Architecture
-	exportImageConfig.Variant = pctx.Platform.Get().Variant
+	img := NewImage(imageConfig, pctx.Platform.Get())
 
 	// Export image
 	lg.Debug().Str("dest", dest.String()).Msg("export image")
-	resp, err := s.Export(ctx, st, exportImageConfig, bk.ExportEntry{
+	resp, err := s.Export(ctx, st, &img, bk.ExportEntry{
 		Type: bk.ExporterImage,
 		Attrs: map[string]string{
 			"name": dest.String(),