Add load subcommand

Author: twz123

client/load.go

@@ -0,0 +1,135 @@
+package client
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/images/archive"
+	"github.com/docker/distribution/reference"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+
+	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+)
+
+// LoadProgress receives informational updates about the progress of image loading.
+type LoadProgress interface {
+	// LoadImageCreated indicates that the given image has been created.
+	LoadImageCreated(images.Image) error
+
+	// LoadImageUpdated indicates that the given image has been updated.
+	LoadImageUpdated(images.Image) error
+
+	// LoadImageReplaced indicates that an image has been replaced, leaving the old image without a name.
+	LoadImageReplaced(before, after images.Image) error
+}
+
+// LoadImage imports an image into the image store.
+func (c *Client) LoadImage(ctx context.Context, reader io.Reader, mon LoadProgress) error {
+	// Create the worker opts.
+	opt, err := c.createWorkerOpt(false)
+	if err != nil {
+		return errors.Wrap(err, "creating worker opt failed")
+	}
+
+	if opt.ImageStore == nil {
+		return errors.New("image store is nil")
+	}
+
+	logrus.Debug("Importing index")
+	index, err := importIndex(ctx, opt.ContentStore, reader)
+	if err != nil {
+		return err
+	}
+
+	logrus.Debug("Extracting manifests")
+	manifests := extractManifests(index)
+
+	for _, imageSkel := range manifests {
+		err := load(ctx, opt.ImageStore, imageSkel, mon)
+		if err != nil {
+			return nil
+		}
+	}
+
+	return nil
+}
+
+func importIndex(ctx context.Context, store content.Store, reader io.Reader) (*ocispec.Index, error) {
+	d, err := archive.ImportIndex(ctx, store, reader)
+	if err != nil {
+		return nil, err
+	}
+
+	indexBytes, err := content.ReadBlob(ctx, store, d)
+	if err != nil {
+		return nil, err
+	}
+
+	var index ocispec.Index
+	err = json.Unmarshal(indexBytes, &index)
+	if err != nil {
+		return nil, err
+	}
+
+	return &index, nil
+}
+
+func extractManifests(index *ocispec.Index) []images.Image {
+	var result []images.Image
+	for _, m := range index.Manifests {
+		switch m.MediaType {
+		case images.MediaTypeDockerSchema2Manifest:
+			if name, ok := m.Annotations[images.AnnotationImageName]; ok {
+				if ref, ok := m.Annotations[ocispec.AnnotationRefName]; ok {
+					if normalized, err := reference.ParseNormalizedNamed(name); err == nil {
+						if normalizedWithTag, err := reference.WithTag(normalized, ref); err == nil {
+							if normalized == normalizedWithTag {
+								result = append(result, images.Image{
+									Name:   normalized.String(),
+									Target: m,
+								})
+								continue
+							}
+						}
+					}
+				}
+			}
+			logrus.Debugf("Failed to extract image info from manifest: %v", m)
+		}
+	}
+
+	return result
+}
+
+func load(ctx context.Context, store images.Store, imageSkel images.Image, mon LoadProgress) error {
+	image, err := store.Get(ctx, imageSkel.Name)
+
+	if errors.Cause(err) == errdefs.ErrNotFound {
+		image, err = store.Create(ctx, imageSkel)
+		if err != nil {
+			return err
+		}
+		return mon.LoadImageCreated(image)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	updated, err := store.Update(ctx, imageSkel)
+	if err != nil {
+		return err
+	}
+
+	if image.Target.Digest == updated.Target.Digest {
+		return mon.LoadImageUpdated(updated)
+	}
+
+	image.Name = ""
+	return mon.LoadImageReplaced(image, updated)
+}

load.go

@@ -0,0 +1,108 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/namespaces"
+	"github.com/genuinetools/img/client"
+	"github.com/moby/buildkit/identity"
+	"github.com/moby/buildkit/session"
+)
+
+// TODO(AkihiroSuda): support OCI archive
+const loadUsageShortHelp = `Load an image from a tar archive or STDIN.`
+const loadUsageLongHelp = `Load an image from a tar archive or STDIN.`
+
+func newLoadCommand() *cobra.Command {
+
+	load := &loadCommand{}
+
+	cmd := &cobra.Command{
+		Use:                   "load [OPTIONS]",
+		DisableFlagsInUseLine: true,
+		SilenceUsage:          true,
+		Short:                 loadUsageShortHelp,
+		Long:                  loadUsageLongHelp,
+		// Args:                  load.ValidateArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return load.Run(args)
+		},
+	}
+
+	fs := cmd.Flags()
+
+	fs.StringVarP(&load.input, "input", "i", "", "Read from tar archive file, instead of STDIN")
+	fs.BoolVarP(&load.quiet, "quiet", "q", false, "Suppress the load output")
+
+	return cmd
+}
+
+type loadCommand struct {
+	input string
+	quiet bool
+}
+
+func (cmd *loadCommand) Run(args []string) (err error) {
+	reexec()
+
+	// Create the context.
+	id := identity.NewID()
+	ctx := session.NewContext(context.Background(), id)
+	ctx = namespaces.WithNamespace(ctx, "buildkit")
+
+	// Create the client.
+	c, err := client.New(stateDir, backend, nil)
+	if err != nil {
+		return err
+	}
+	defer c.Close()
+
+	// Create the reader.
+	reader, err := cmd.reader()
+	if err != nil {
+		return err
+	}
+	defer reader.Close()
+
+	// Load images.
+	return c.LoadImage(ctx, reader, cmd)
+}
+
+func (cmd *loadCommand) reader() (io.ReadCloser, error) {
+	if cmd.input != "" {
+		return os.Open(cmd.input)
+	}
+
+	return os.Stdin, nil
+}
+
+func (cmd *loadCommand) LoadImageCreated(image images.Image) error {
+	if cmd.quiet {
+		return nil
+	}
+	_, err := fmt.Printf("Loaded image: %s\n", image.Name)
+	return err
+}
+
+func (cmd *loadCommand) LoadImageUpdated(image images.Image) error {
+	return cmd.LoadImageCreated(image)
+}
+
+func (cmd *loadCommand) LoadImageReplaced(before, after images.Image) error {
+	if !cmd.quiet {
+		_, err := fmt.Printf(
+			"The image %s already exists, leaving the old one with ID %s orphaned\n",
+			after.Name, before.Target.Digest)
+		if err != nil {
+			return err
+		}
+	}
+
+	return cmd.LoadImageCreated(after)
+}

load_test.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestLoadImage(t *testing.T) {
+	output := run(t, "load", "-i", "testdata/img-load/oci.tar")
+	if output != "Loaded image: docker.io/library/dummy:latest\n" {
+		t.Fatalf("Unexpected output: %s", output)
+	}
+
+	output = run(t, "load", "-i", "testdata/img-load/oci.tar")
+	if output != "Loaded image: docker.io/library/dummy:latest\n" {
+		t.Fatalf("Unexpected output: %s", output)
+	}
+
+	output = run(t, "load", "-i", "testdata/img-load/docker.tar")
+	expected := `The image docker.io/library/dummy:latest already exists, leaving the old one with ID sha256:e08488191147b6fc575452dfac3238721aa5b86d2545a9edaa1c1d88632b2233 orphaned
+Loaded image: docker.io/library/dummy:latest
+`
+	if output != expected {
+		t.Fatalf("Unexpected output: %s", output)
+	}
+
+	output = run(t, "load", "-i", "testdata/img-load/docker.tar")
+	if output != "Loaded image: docker.io/library/dummy:latest\n" {
+		t.Fatalf("Unexpected output: %s", output)
+	}
+}

main.go

@@ -75,6 +75,7 @@ func main() {
 		newDiskUsageCommand(),
 		newInspectCommand(),
 		newListCommand(),
+		newLoadCommand(),
 		newLoginCommand(),
 		newLogoutCommand(),
 		newPruneCommand(),