diff options
| -rw-r--r-- | drive/list.go | 56 | ||||
| -rw-r--r-- | drive/sync.go | 18 | ||||
| -rw-r--r-- | drive/util.go | 5 |
3 files changed, 58 insertions, 21 deletions
diff --git a/drive/list.go b/drive/list.go index d3d8f3d..e636585 100644 --- a/drive/list.go +++ b/drive/list.go @@ -4,7 +4,9 @@ import ( "fmt" "io" "text/tabwriter" + "golang.org/x/net/context" "google.golang.org/api/drive/v3" + "google.golang.org/api/googleapi" ) type ListFilesArgs struct { @@ -18,14 +20,20 @@ type ListFilesArgs struct { } func (self *Drive) List(args ListFilesArgs) (err error) { - fileList, err := self.service.Files.List().PageSize(args.MaxFiles).Q(args.Query).OrderBy(args.SortOrder).Fields("files(id,name,md5Checksum,mimeType,size,createdTime)").Do() + listArgs := listAllFilesArgs{ + query: args.Query, + fields: []googleapi.Field{"nextPageToken", "files(id,name,md5Checksum,mimeType,size,createdTime)"}, + sortOrder: args.SortOrder, + maxFiles: args.MaxFiles, + } + files, err := self.listAllFiles(listArgs) if err != nil { - return fmt.Errorf("Failed listing files: %s", err) + return fmt.Errorf("Failed to list files: %s", err) } PrintFileList(PrintFileListArgs{ Out: args.Out, - Files: fileList.Files, + Files: files, NameWidth: int(args.NameWidth), SkipHeader: args.SkipHeader, SizeInBytes: args.SizeInBytes, @@ -34,6 +42,48 @@ func (self *Drive) List(args ListFilesArgs) (err error) { return } +type listAllFilesArgs struct { + query string + fields []googleapi.Field + sortOrder string + maxFiles int64 +} + +func (self *Drive) listAllFiles(args listAllFilesArgs) ([]*drive.File, error) { + var files []*drive.File + + var pageSize int64 + if args.maxFiles > 0 && args.maxFiles < 1000 { + pageSize = args.maxFiles + } else { + pageSize = 1000 + } + + controlledStop := fmt.Errorf("Controlled stop") + + err := self.service.Files.List().Q(args.query).Fields(args.fields...).OrderBy(args.sortOrder).PageSize(pageSize).Pages(context.TODO(), func(fl *drive.FileList) error { + files = append(files, fl.Files...) + + // Stop when we have all the files we need + if args.maxFiles > 0 && len(files) >= int(args.maxFiles) { + return controlledStop + } + + return nil + }) + + if err != nil && err != controlledStop { + return nil, err + } + + if args.maxFiles > 0 { + n := min(len(files), int(args.maxFiles)) + return files[:n], nil + } + + return files, nil +} + type PrintFileListArgs struct { Out io.Writer Files []*drive.File diff --git a/drive/sync.go b/drive/sync.go index c0d5add..86ddb3b 100644 --- a/drive/sync.go +++ b/drive/sync.go @@ -8,7 +8,6 @@ import ( "path/filepath" "github.com/soniakeys/graph" "github.com/sabhiram/go-git-ignore" - "golang.org/x/net/context" "google.golang.org/api/drive/v3" "google.golang.org/api/googleapi" ) @@ -104,23 +103,6 @@ func prepareLocalFiles(root string) ([]*LocalFile, error) { return files, err } -type listAllFilesArgs struct { - query string - fields []googleapi.Field - sortOrder string -} - -func (self *Drive) listAllFiles(args listAllFilesArgs) ([]*drive.File, error) { - var files []*drive.File - - err := self.service.Files.List().Q(args.query).Fields(args.fields...).OrderBy(args.sortOrder).PageSize(1000).Pages(context.TODO(), func(fl *drive.FileList) error { - files = append(files, fl.Files...) - return nil - }) - - return files, err -} - func (self *Drive) prepareRemoteFiles(rootDir *drive.File, sortOrder string) ([]*RemoteFile, error) { // Find all files which has rootDir as root listArgs := listAllFilesArgs{ diff --git a/drive/util.go b/drive/util.go index 481a48e..7a7845a 100644 --- a/drive/util.go +++ b/drive/util.go @@ -148,3 +148,8 @@ func pow(x int, y int) int { f := math.Pow(float64(x), float64(y)) return int(f) } + +func min(x int, y int) int { + n := math.Min(float64(x), float64(y)) + return int(n) +} |