Skip to content

Commit

Permalink
refactor: create tty copy handler
Browse files Browse the repository at this point in the history
Signed-off-by: Terry Howe <terrylhowe@gmail.com>
  • Loading branch information
TerryHowe committed Aug 22, 2024
1 parent 5792c35 commit 9676775
Show file tree
Hide file tree
Showing 8 changed files with 222 additions and 77 deletions.
5 changes: 4 additions & 1 deletion cmd/oras/internal/display/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,9 @@ func NewManifestPushHandler(printer *output.Printer) metadata.ManifestPushHandle
}

// NewCopyHandler returns copy handlers.
func NewCopyHandler(printer *output.Printer, fetcher fetcher.Fetcher) (status.CopyHandler, metadata.CopyHandler) {
func NewCopyHandler(printer *output.Printer, tty *os.File, fetcher fetcher.Fetcher) (status.CopyHandler, metadata.CopyHandler) {
if tty != nil {
return status.NewTTYCopyHandler(tty), text.NewCopyHandler(printer)
}
return status.NewTextCopyHandler(printer, fetcher), text.NewCopyHandler(printer)
}
23 changes: 22 additions & 1 deletion cmd/oras/internal/display/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,15 @@ limitations under the License.
package display

import (
"oras.land/oras/internal/testutils"
"oras.land/oras/cmd/oras/internal/display/metadata/text"
"os"
"reflect"
"testing"

"oras.land/oras/cmd/oras/internal/display/status"
"oras.land/oras/cmd/oras/internal/option"
"oras.land/oras/cmd/oras/internal/output"
"oras.land/oras/internal/testutils"
)

func TestNewPushHandler(t *testing.T) {
Expand Down Expand Up @@ -49,3 +52,21 @@ func TestNewPullHandler(t *testing.T) {
t.Errorf("NewPullHandler() error = %v, want nil", err)
}
}

func TestNewCopyHandler(t *testing.T) {
printer := output.NewPrinter(os.Stdout, os.Stderr, false)
copyHandler, copyMetadataHandler := NewCopyHandler(printer, os.Stdout, nil)
if _, ok := copyHandler.(*status.TTYCopyHandler); !ok {
t.Errorf("expected *status.TTYCopyHandler actual %v", reflect.TypeOf(copyHandler))
}
if _, ok := copyMetadataHandler.(*text.CopyHandler); !ok {
t.Errorf("expected metadata.CopyHandler actual %v", reflect.TypeOf(copyMetadataHandler))
}
copyHandler, copyMetadataHandler = NewCopyHandler(printer, nil, nil)
if _, ok := copyHandler.(*status.TextCopyHandler); !ok {
t.Errorf("expected *status.TextCopyHandler actual %v", reflect.TypeOf(copyHandler))
}
if _, ok := copyMetadataHandler.(*text.CopyHandler); !ok {
t.Errorf("expected metadata.CopyHandler actual %v", reflect.TypeOf(copyMetadataHandler))
}
}
2 changes: 2 additions & 0 deletions cmd/oras/internal/display/status/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,4 +60,6 @@ type CopyHandler interface {
PreCopy(ctx context.Context, desc ocispec.Descriptor) error
PostCopy(ctx context.Context, desc ocispec.Descriptor) error
OnMounted(ctx context.Context, desc ocispec.Descriptor) error
StartTracking(gt oras.GraphTarget) (oras.GraphTarget, error)
StopTracking()
}
9 changes: 9 additions & 0 deletions cmd/oras/internal/display/status/text.go
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,15 @@ func NewTextCopyHandler(printer *output.Printer, fetcher content.Fetcher) CopyHa
}
}

// StartTracking starts a tracked target from a graph target.
func (ch *TextCopyHandler) StartTracking(gt oras.GraphTarget) (oras.GraphTarget, error) {
return gt, nil
}

// StopTracking ends the copy tracking for the target.
func (ch *TextCopyHandler) StopTracking() {
}

// OnCopySkipped is called when an object already exists.
func (ch *TextCopyHandler) OnCopySkipped(_ context.Context, desc ocispec.Descriptor) error {
ch.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
Expand Down
59 changes: 59 additions & 0 deletions cmd/oras/internal/display/status/tty.go
Original file line number Diff line number Diff line change
Expand Up @@ -143,3 +143,62 @@ func (ph *TTYPullHandler) TrackTarget(gt oras.GraphTarget) (oras.GraphTarget, St
ph.tracked = tracked
return tracked, tracked.Close, nil
}

// TTYCopyHandler handles tty status output for copy events.
type TTYCopyHandler struct {
tty *os.File
committed *sync.Map
tracked track.GraphTarget
}

// NewTTYCopyHandler returns a new handler for copy command.
func NewTTYCopyHandler(tty *os.File) CopyHandler {
return &TTYCopyHandler{
tty: tty,
committed: &sync.Map{},
}
}

// StartTracking returns a tracked target from a graph target.
func (ch *TTYCopyHandler) StartTracking(gt oras.GraphTarget) (oras.GraphTarget, error) {
tracked, err := track.NewTarget(gt, copyPromptCopying, copyPromptCopied, ch.tty)
ch.tracked = tracked
return tracked, err
}

// StopTracking ends the copy tracking for the target.
func (ch *TTYCopyHandler) StopTracking() {
_ = ch.tracked.Close()
}

// OnCopySkipped is called when an object already exists.
func (ch *TTYCopyHandler) OnCopySkipped(_ context.Context, desc ocispec.Descriptor) error {
ch.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return ch.tracked.Prompt(desc, copyPromptExists)
}

// PreCopy implements PreCopy of CopyHandler.
func (ch *TTYCopyHandler) PreCopy(_ context.Context, _ ocispec.Descriptor) error {
return nil
}

// PostCopy implements PostCopy of CopyHandler.
func (ch *TTYCopyHandler) PostCopy(ctx context.Context, desc ocispec.Descriptor) error {
ch.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
successors, err := graph.FilteredSuccessors(ctx, desc, ch.tracked, DeduplicatedFilter(ch.committed))
if err != nil {
return err
}
for _, successor := range successors {
if err = ch.tracked.Prompt(successor, copyPromptSkipped); err != nil {
return err

Check warning on line 194 in cmd/oras/internal/display/status/tty.go

View check run for this annotation

Codecov / codecov/patch

cmd/oras/internal/display/status/tty.go#L193-L194

Added lines #L193 - L194 were not covered by tests
}
}
return nil
}

// OnMounted implements OnMounted of CopyHandler.
func (ch *TTYCopyHandler) OnMounted(_ context.Context, desc ocispec.Descriptor) error {
ch.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return ch.tracked.Prompt(desc, copyPromptMounted)
}
93 changes: 93 additions & 0 deletions cmd/oras/internal/display/status/tty_console_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,16 @@ limitations under the License.
package status

import (
"oras.land/oras-go/v2"
"oras.land/oras-go/v2/content/memory"
"oras.land/oras/internal/testutils"
"testing"
)

type testGraphTarget struct {
oras.GraphTarget
}

func TestTTYPushHandler_TrackTarget(t *testing.T) {
// prepare pty
_, slave, err := testutils.NewPty()
Expand Down Expand Up @@ -78,3 +83,91 @@ func Test_TTYPullHandler_TrackTarget(t *testing.T) {
}
})
}

func TestTTYCopyHandler_OnMounted(t *testing.T) {
pty, slave, err := testutils.NewPty()
if err != nil {
t.Fatal(err)
}
defer slave.Close()
ch := NewTTYCopyHandler(slave)
_, err = ch.StartTracking(&testGraphTarget{memory.New()})
if err != nil {
t.Fatal(err)
}
defer ch.StopTracking()

if err = ch.OnMounted(ctx, mockFetcher.OciImage); err != nil {
t.Errorf("OnMounted() should not return an error: %v", err)
}

if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
t.Fatal(err)
}
}

func TestTTYCopyHandler_OnCopySkipped(t *testing.T) {
pty, slave, err := testutils.NewPty()
if err != nil {
t.Fatal(err)
}
defer slave.Close()
ch := NewTTYCopyHandler(slave)
_, err = ch.StartTracking(&testGraphTarget{memory.New()})
if err != nil {
t.Fatal(err)
}
defer ch.StopTracking()

if err = ch.OnCopySkipped(ctx, mockFetcher.OciImage); err != nil {
t.Errorf("OnCopySkipped() should not return an error: %v", err)
}

if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
t.Fatal(err)
}
}

func TestTTYCopyHandler_PostCopy(t *testing.T) {
pty, slave, err := testutils.NewPty()
if err != nil {
t.Fatal(err)
}
defer slave.Close()
ch := NewTTYCopyHandler(slave)
_, err = ch.StartTracking(&testGraphTarget{memory.New()})
if err != nil {
t.Fatal(err)
}
defer ch.StopTracking()

if ch.PostCopy(ctx, bogus) == nil {
t.Error("PostCopy() should return an error")
}

if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
t.Fatal(err)
}
}

func TestTTYCopyHandler_PreCopy(t *testing.T) {
pty, slave, err := testutils.NewPty()
if err != nil {
t.Fatal(err)
}
defer slave.Close()
ch := NewTTYCopyHandler(slave)
_, err = ch.StartTracking(&testGraphTarget{memory.New()})
if err != nil {
t.Fatal(err)
}
defer ch.StopTracking()

if err = ch.PreCopy(ctx, mockFetcher.OciImage); err != nil {
t.Errorf("PreCopy() should not return an error: %v", err)
}

if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
t.Fatal(err)
}
}
61 changes: 11 additions & 50 deletions cmd/oras/root/cp.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,8 @@ import (
"context"
"encoding/json"
"fmt"
"oras.land/oras/cmd/oras/internal/display/status"
"slices"
"strings"
"sync"

"github.com/opencontainers/go-digest"
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
Expand All @@ -35,7 +33,7 @@ import (
"oras.land/oras/cmd/oras/internal/argument"
"oras.land/oras/cmd/oras/internal/command"
"oras.land/oras/cmd/oras/internal/display"
"oras.land/oras/cmd/oras/internal/display/status/track"
"oras.land/oras/cmd/oras/internal/display/status"
oerrors "oras.land/oras/cmd/oras/internal/errors"
"oras.land/oras/cmd/oras/internal/option"
"oras.land/oras/internal/docker"
Expand Down Expand Up @@ -127,7 +125,7 @@ func runCopy(cmd *cobra.Command, opts *copyOptions) error {
return err
}
ctx = registryutil.WithScopeHint(ctx, dst, auth.ActionPull, auth.ActionPush)
copyHandler, handler := display.NewCopyHandler(opts.Printer, dst)
copyHandler, handler := display.NewCopyHandler(opts.Printer, opts.TTY, dst)

desc, err := doCopy(ctx, copyHandler, src, dst, opts)
if err != nil {
Expand All @@ -154,68 +152,31 @@ func runCopy(cmd *cobra.Command, opts *copyOptions) error {
return nil
}

func doCopy(ctx context.Context, copyHandler status.CopyHandler, src oras.ReadOnlyGraphTarget, dst oras.GraphTarget, opts *copyOptions) (ocispec.Descriptor, error) {
func doCopy(ctx context.Context, copyHandler status.CopyHandler, src oras.ReadOnlyGraphTarget, dst oras.GraphTarget, opts *copyOptions) (desc ocispec.Descriptor, err error) {
// Prepare copy options
committed := &sync.Map{}
extendedCopyOptions := oras.DefaultExtendedCopyOptions
extendedCopyOptions.Concurrency = opts.concurrency
extendedCopyOptions.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
return registry.Referrers(ctx, src, desc, "")
}

const (
promptExists = "Exists "
promptCopying = "Copying"
promptCopied = "Copied "
promptSkipped = "Skipped"
promptMounted = "Mounted"
)
srcRepo, srcIsRemote := src.(*remote.Repository)
dstRepo, dstIsRemote := dst.(*remote.Repository)
if srcIsRemote && dstIsRemote && srcRepo.Reference.Registry == dstRepo.Reference.Registry {
extendedCopyOptions.MountFrom = func(ctx context.Context, desc ocispec.Descriptor) ([]string, error) {
return []string{srcRepo.Reference.Repository}, nil
}
}
if opts.TTY == nil {
// no TTY output
extendedCopyOptions.OnCopySkipped = copyHandler.OnCopySkipped
extendedCopyOptions.PreCopy = copyHandler.PreCopy
extendedCopyOptions.PostCopy = copyHandler.PostCopy
extendedCopyOptions.OnMounted = copyHandler.OnMounted
} else {
// TTY output
tracked, err := track.NewTarget(dst, promptCopying, promptCopied, opts.TTY)
if err != nil {
return ocispec.Descriptor{}, err
}
defer tracked.Close()
dst = tracked
extendedCopyOptions.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return tracked.Prompt(desc, promptExists)
}
extendedCopyOptions.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
successors, err := graph.FilteredSuccessors(ctx, desc, tracked, status.DeduplicatedFilter(committed))
if err != nil {
return err
}
for _, successor := range successors {
if err = tracked.Prompt(successor, promptSkipped); err != nil {
return err
}
}
return nil
}
extendedCopyOptions.OnMounted = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return tracked.Prompt(desc, promptMounted)
}
dst, err = copyHandler.StartTracking(dst)
if err != nil {
return desc, err

Check warning on line 172 in cmd/oras/root/cp.go

View check run for this annotation

Codecov / codecov/patch

cmd/oras/root/cp.go#L172

Added line #L172 was not covered by tests
}
defer copyHandler.StopTracking()
extendedCopyOptions.OnCopySkipped = copyHandler.OnCopySkipped
extendedCopyOptions.PreCopy = copyHandler.PreCopy
extendedCopyOptions.PostCopy = copyHandler.PostCopy
extendedCopyOptions.OnMounted = copyHandler.OnMounted

var desc ocispec.Descriptor
var err error
rOpts := oras.DefaultResolveOptions
rOpts.TargetPlatform = opts.Platform.Platform
if opts.recursive {
Expand Down
Loading

0 comments on commit 9676775

Please sign in to comment.