support partial download on get command

Signed-off-by: Wei-Chun, Chang <wcchang@infuseai.io>
InfuseAI · Mar 22, 2022 · 4260d84 · 4260d84
1 parent 3ee907e
commit 4260d84
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 15 deletions.
diff --git a/cmd/get.go b/cmd/get.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"errors"
 	"os"
 	"path/filepath"
 	"strings"
@@ -22,11 +23,15 @@ var getCmd = &cobra.Command{
   
   # Download to a specific folder
   art get -o /tmp/mydataset s3://bucket/mydataset`,
-	Args: cobra.ExactArgs(1),
+	Args: cobra.MinimumNArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
 		var err error
 
 		repoUrl, ref, err := parseRepoStr(args[0])
+		if err != nil {
+			exitWithError(err)
+			return
+		}
 		baseDir, err := cmd.Flags().GetString("output")
 		if err != nil {
 			exitWithError(err)
@@ -68,6 +73,17 @@ var getCmd = &cobra.Command{
 			exitWithError(err)
 		}
 
+		if len(args) > 1 {
+			if options.Delete {
+				exitWithError(errors.New("cannot download partial files and specify delete flag at the same time"))
+			}
+			fileInclude := core.NewAvcInclude(args[1:])
+			fileFilter := func(path string) bool {
+				return fileInclude.MatchesPath(path)
+			}
+			options.FileFilter = fileFilter
+		}
+
 		err = mngr.Pull(options)
 		if err != nil {
 			exitWithError(err)

diff --git a/internal/core/config.go b/internal/core/config.go
@@ -213,3 +213,11 @@ func NewAvcIgnore(dir string) (*AvcIgnore, error) {
 
 	return avcIgnore, nil
 }
+
+type AvcInclude = gitignore.GitIgnore
+
+func NewAvcInclude(filePath []string) *AvcInclude {
+	filter := gitignore.CompileIgnoreLines(filePath...)
+
+	return filter
+}
diff --git a/internal/core/manager.go b/internal/core/manager.go
@@ -610,12 +610,13 @@ func (mngr *ArtifactManager) Pull(options PullOptions) error {
 
 	// Diff
 	result, err := mngr.Diff(DiffOptions{
-		NoDelete:     !options.Delete,
-		LeftCommit:   commitLocal,
-		RightCommit:  commitRemote,
-		AddFilter:    avcIgnoreFilter,
-		ChangeFilter: avcIgnoreFilter,
-		DeleteFilter: avcIgnoreFilter,
+		NoDelete:      !options.Delete,
+		LeftCommit:    commitLocal,
+		RightCommit:   commitRemote,
+		AddFilter:     avcIgnoreFilter,
+		ChangeFilter:  avcIgnoreFilter,
+		DeleteFilter:  avcIgnoreFilter,
+		IncludeFilter: options.FileFilter,
 	})
 	if err != nil {
 		return err
@@ -817,6 +818,9 @@ func (mngr *ArtifactManager) Diff(option DiffOptions) (DiffResult, error) {
 		}
 	}
 	for i, blob := range leftCommit.Blobs {
+		if option.IncludeFilter != nil && !option.IncludeFilter(blob.Path) {
+			continue
+		}
 		entry := entries[blob.Path]
 		entry.left = &leftCommit.Blobs[i]
 		entries[blob.Path] = entry
@@ -837,6 +841,9 @@ func (mngr *ArtifactManager) Diff(option DiffOptions) (DiffResult, error) {
 	}
 
 	for i, blob := range rightCommit.Blobs {
+		if option.IncludeFilter != nil && !option.IncludeFilter(blob.Path) {
+			continue
+		}
 		entry := entries[blob.Path]
 		entry.right = &rightCommit.Blobs[i]
 		entries[blob.Path] = entry

diff --git a/internal/core/types.go b/internal/core/types.go
@@ -40,19 +40,21 @@ type PullOptions struct {
 	NoFetch     bool
 	Delete      bool
 	RefOrCommit *string
+	FileFilter  PathFilter
 }
 
 type PathFilter func(path string) bool
 
 type DiffOptions struct {
-	LeftRef      string
-	LeftCommit   *Commit
-	RightRef     string
-	RightCommit  *Commit
-	AddFilter    PathFilter
-	ChangeFilter PathFilter
-	DeleteFilter PathFilter
-	NoDelete     bool
+	LeftRef       string
+	LeftCommit    *Commit
+	RightRef      string
+	RightCommit   *Commit
+	AddFilter     PathFilter
+	ChangeFilter  PathFilter
+	DeleteFilter  PathFilter
+	IncludeFilter PathFilter
+	NoDelete      bool
 }
 
 type DiffType int