Support listing more than 1000 files

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)
+}