refactor: Successor getting with separation of concerns (oras-project…

…#1443) Signed-off-by: Terry Howe <terrylhowe@gmail.com> Signed-off-by: Billy Zha <jinzha1@microsoft.com> Co-authored-by: Billy Zha <jinzha1@microsoft.com>
qweeah · Aug 9, 2024 · 141c5e1 · 141c5e1
1 parent c31d120
commit 141c5e1
Show file tree

Hide file tree

Showing 20 changed files with 588 additions and 222 deletions.
diff --git a/cmd/oras/internal/display/handler.go b/cmd/oras/internal/display/handler.go
@@ -37,12 +37,12 @@ import (
 )
 
 // NewPushHandler returns status and metadata handlers for push command.
-func NewPushHandler(printer *output.Printer, format option.Format, tty *os.File) (status.PushHandler, metadata.PushHandler, error) {
+func NewPushHandler(printer *output.Printer, format option.Format, tty *os.File, fetcher fetcher.Fetcher) (status.PushHandler, metadata.PushHandler, error) {
 	var statusHandler status.PushHandler
 	if tty != nil {
-		statusHandler = status.NewTTYPushHandler(tty)
+		statusHandler = status.NewTTYPushHandler(tty, fetcher)
 	} else if format.Type == option.FormatTypeText.Name {
-		statusHandler = status.NewTextPushHandler(printer)
+		statusHandler = status.NewTextPushHandler(printer, fetcher)
 	} else {
 		statusHandler = status.NewDiscardHandler()
 	}
@@ -62,12 +62,12 @@ func NewPushHandler(printer *output.Printer, format option.Format, tty *os.File)
 }
 
 // NewAttachHandler returns status and metadata handlers for attach command.
-func NewAttachHandler(printer *output.Printer, format option.Format, tty *os.File) (status.AttachHandler, metadata.AttachHandler, error) {
+func NewAttachHandler(printer *output.Printer, format option.Format, tty *os.File, fetcher fetcher.Fetcher) (status.AttachHandler, metadata.AttachHandler, error) {
 	var statusHandler status.AttachHandler
 	if tty != nil {
-		statusHandler = status.NewTTYAttachHandler(tty)
+		statusHandler = status.NewTTYAttachHandler(tty, fetcher)
 	} else if format.Type == option.FormatTypeText.Name {
-		statusHandler = status.NewTextAttachHandler(printer)
+		statusHandler = status.NewTextAttachHandler(printer, fetcher)
 	} else {
 		statusHandler = status.NewDiscardHandler()
 	}

diff --git a/cmd/oras/internal/display/handler_test.go b/cmd/oras/internal/display/handler_test.go
@@ -16,6 +16,7 @@ limitations under the License.
 package display
 
 import (
+	"oras.land/oras/internal/testutils"
 	"os"
 	"testing"
 
@@ -24,16 +25,18 @@ import (
 )
 
 func TestNewPushHandler(t *testing.T) {
+	mockFetcher := testutils.NewMockFetcher()
 	printer := output.NewPrinter(os.Stdout, os.Stderr, false)
-	_, _, err := NewPushHandler(printer, option.Format{Type: option.FormatTypeText.Name}, os.Stdout)
+	_, _, err := NewPushHandler(printer, option.Format{Type: option.FormatTypeText.Name}, os.Stdout, mockFetcher.Fetcher)
 	if err != nil {
 		t.Errorf("NewPushHandler() error = %v, want nil", err)
 	}
 }
 
 func TestNewAttachHandler(t *testing.T) {
+	mockFetcher := testutils.NewMockFetcher()
 	printer := output.NewPrinter(os.Stdout, os.Stderr, false)
-	_, _, err := NewAttachHandler(printer, option.Format{Type: option.FormatTypeText.Name}, os.Stdout)
+	_, _, err := NewAttachHandler(printer, option.Format{Type: option.FormatTypeText.Name}, os.Stdout, mockFetcher.Fetcher)
 	if err != nil {
 		t.Errorf("NewAttachHandler() error = %v, want nil", err)
 	}

diff --git a/cmd/oras/internal/display/status/discard.go b/cmd/oras/internal/display/status/discard.go
@@ -16,9 +16,10 @@ limitations under the License.
 package status
 
 import (
+	"context"
+
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2"
-	"oras.land/oras-go/v2/content"
 )
 
 func discardStopTrack() error {
@@ -48,8 +49,20 @@ func (DiscardHandler) TrackTarget(gt oras.GraphTarget) (oras.GraphTarget, StopTr
 	return gt, discardStopTrack, nil
 }
 
-// UpdateCopyOptions updates the copy options for the artifact push.
-func (DiscardHandler) UpdateCopyOptions(opts *oras.CopyGraphOptions, fetcher content.Fetcher) {}
+// OnCopySkipped is called when an object already exists.
+func (DiscardHandler) OnCopySkipped(_ context.Context, _ ocispec.Descriptor) error {
+	return nil
+}
+
+// PreCopy implements PreCopy of CopyHandler.
+func (DiscardHandler) PreCopy(_ context.Context, _ ocispec.Descriptor) error {
+	return nil
+}
+
+// PostCopy implements PostCopy of CopyHandler.
+func (DiscardHandler) PostCopy(_ context.Context, _ ocispec.Descriptor) error {
+	return nil
+}
 
 // OnNodeDownloading implements PullHandler.
 func (DiscardHandler) OnNodeDownloading(desc ocispec.Descriptor) error {

diff --git a/cmd/oras/internal/display/status/interface.go b/cmd/oras/internal/display/status/interface.go
@@ -19,7 +19,6 @@ import (
 	"context"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2"
-	"oras.land/oras-go/v2/content"
 )
 
 // StopTrackTargetFunc is the function type to stop tracking a target.
@@ -30,7 +29,9 @@ type PushHandler interface {
 	OnFileLoading(name string) error
 	OnEmptyArtifact() error
 	TrackTarget(gt oras.GraphTarget) (oras.GraphTarget, StopTrackTargetFunc, error)
-	UpdateCopyOptions(opts *oras.CopyGraphOptions, fetcher content.Fetcher)
+	OnCopySkipped(ctx context.Context, desc ocispec.Descriptor) error
+	PreCopy(ctx context.Context, desc ocispec.Descriptor) error
+	PostCopy(ctx context.Context, desc ocispec.Descriptor) error
 }
 
 // AttachHandler handles text status output for attach command.

diff --git a/cmd/oras/internal/display/status/text.go b/cmd/oras/internal/display/status/text.go
@@ -17,25 +17,30 @@ package status
 
 import (
 	"context"
+	"oras.land/oras/internal/graph"
 	"sync"
 
-	"oras.land/oras/cmd/oras/internal/output"
-
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"oras.land/oras-go/v2"
 	"oras.land/oras-go/v2/content"
+	"oras.land/oras/cmd/oras/internal/output"
 )
 
 // TextPushHandler handles text status output for push events.
 type TextPushHandler struct {
-	printer *output.Printer
+	printer   *output.Printer
+	committed *sync.Map
+	fetcher   content.Fetcher
 }
 
 // NewTextPushHandler returns a new handler for push command.
-func NewTextPushHandler(printer *output.Printer) PushHandler {
-	return &TextPushHandler{
-		printer: printer,
+func NewTextPushHandler(printer *output.Printer, fetcher content.Fetcher) PushHandler {
+	tch := TextPushHandler{
+		printer:   printer,
+		fetcher:   fetcher,
+		committed: &sync.Map{},
 	}
+	return &tch
 }
 
 // OnFileLoading is called when a file is being prepared for upload.
@@ -53,28 +58,35 @@ func (ph *TextPushHandler) TrackTarget(gt oras.GraphTarget) (oras.GraphTarget, S
 	return gt, discardStopTrack, nil
 }
 
-// UpdateCopyOptions adds status update to the copy options.
-func (ph *TextPushHandler) UpdateCopyOptions(opts *oras.CopyGraphOptions, fetcher content.Fetcher) {
-	committed := &sync.Map{}
-	opts.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error {
-		committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
-		return ph.printer.PrintStatus(desc, PushPromptExists)
-	}
-	opts.PreCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
-		return ph.printer.PrintStatus(desc, PushPromptUploading)
+// OnCopySkipped is called when an object already exists.
+func (ph *TextPushHandler) OnCopySkipped(_ context.Context, desc ocispec.Descriptor) error {
+	ph.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
+	return ph.printer.PrintStatus(desc, PushPromptExists)
+}
+
+// PreCopy implements PreCopy of CopyHandler.
+func (ph *TextPushHandler) PreCopy(_ context.Context, desc ocispec.Descriptor) error {
+	return ph.printer.PrintStatus(desc, PushPromptUploading)
+}
+
+// PostCopy implements PostCopy of CopyHandler.
+func (ph *TextPushHandler) PostCopy(ctx context.Context, desc ocispec.Descriptor) error {
+	ph.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
+	successors, err := graph.FilteredSuccessors(ctx, desc, ph.fetcher, DeduplicatedFilter(ph.committed))
+	if err != nil {
+		return err
 	}
-	opts.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
-		committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
-		if err := output.PrintSuccessorStatus(ctx, desc, fetcher, committed, ph.printer.StatusPrinter(PushPromptSkipped)); err != nil {
+	for _, successor := range successors {
+		if err = ph.printer.PrintStatus(successor, PushPromptExists); err != nil {
 			return err
 		}
-		return ph.printer.PrintStatus(desc, PushPromptUploaded)
 	}
+	return ph.printer.PrintStatus(desc, PushPromptUploaded)
 }
 
 // NewTextAttachHandler returns a new handler for attach command.
-func NewTextAttachHandler(printer *output.Printer) AttachHandler {
-	return NewTextPushHandler(printer)
+func NewTextAttachHandler(printer *output.Printer, fetcher content.Fetcher) AttachHandler {
+	return NewTextPushHandler(printer, fetcher)
 }
 
 // TextPullHandler handles text status output for pull events.
@@ -149,9 +161,15 @@ func (ch *TextCopyHandler) PreCopy(_ context.Context, desc ocispec.Descriptor) e
 // PostCopy implements PostCopy of CopyHandler.
 func (ch *TextCopyHandler) PostCopy(ctx context.Context, desc ocispec.Descriptor) error {
 	ch.committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
-	if err := output.PrintSuccessorStatus(ctx, desc, ch.fetcher, ch.committed, ch.printer.StatusPrinter(copyPromptSkipped)); err != nil {
+	successors, err := graph.FilteredSuccessors(ctx, desc, ch.fetcher, DeduplicatedFilter(ch.committed))
+	if err != nil {
 		return err
 	}
+	for _, successor := range successors {
+		if err = ch.printer.PrintStatus(successor, copyPromptExists); err != nil {
+			return err
+		}
+	}
 	return ch.printer.PrintStatus(desc, copyPromptCopied)
 }