Enable progress indicator

7 files changed, 60 insertions, 15 deletions

diff --git a/drive/download.go b/drive/download.go
index 31bcc58..e5674ad 100644
--- a/drive/download.go
+++ b/drive/download.go
@@ -8,9 +8,9 @@ import (
 
 type DownloadFileArgs struct {
     Out io.Writer
+    Progress io.Writer
     Id string
     Force bool
-    NoProgress bool
     Stdout bool
 }
 
@@ -30,9 +30,12 @@ func (self *Drive) Download(args DownloadFileArgs) (err error) {
     // Close body on function exit
     defer res.Body.Close()
 
+    // Wrap response body in progress reader
+    srcReader := getProgressReader(res.Body, args.Progress, res.ContentLength)
+
     if args.Stdout {
         // Write file content to stdout
-        _, err := io.Copy(os.Stdout, res.Body)
+        _, err := io.Copy(args.Out, srcReader)
         return err
     }
 
@@ -51,7 +54,7 @@ func (self *Drive) Download(args DownloadFileArgs) (err error) {
     defer outFile.Close()
 
     // Save file to disk
-    bytes, err := io.Copy(outFile, res.Body)
+    bytes, err := io.Copy(outFile, srcReader)
     if err != nil {
         return fmt.Errorf("Failed saving file: %s", err)
     }
diff --git a/drive/progress.go b/drive/progress.go
index ee63b6c..6187058 100644
--- a/drive/progress.go
+++ b/drive/progress.go
@@ -2,10 +2,23 @@ package drive
 
 import (
     "io"
+    "io/ioutil"
     "fmt"
     "time"
 )
 
+func getProgressReader(r io.Reader, w io.Writer, size int64) io.Reader {
+    if w == ioutil.Discard || size < 1024 * 1024 {
+        return r
+    }
+
+    return &Progress{
+        Reader: r,
+        Writer: w,
+        Size: size,
+    }
+}
+
 type Progress struct {
     Writer io.Writer
     Reader io.Reader
diff --git a/drive/revision_download.go b/drive/revision_download.go
index c26c3ce..f06dac3 100644
--- a/drive/revision_download.go
+++ b/drive/revision_download.go
@@ -8,10 +8,10 @@ import (
 
 type DownloadRevisionArgs struct {
     Out io.Writer
+    Progress io.Writer
     FileId string
     RevisionId string
     Force bool
-    NoProgress bool
     Stdout bool
 }
 
@@ -35,9 +35,12 @@ func (self *Drive) DownloadRevision(args DownloadRevisionArgs) (err error) {
     // Close body on function exit
     defer res.Body.Close()
 
+    // Wrap response body in progress reader
+    srcReader := getProgressReader(res.Body, args.Progress, res.ContentLength)
+
     if args.Stdout {
         // Write file content to stdout
-        _, err := io.Copy(os.Stdout, res.Body)
+        _, err := io.Copy(args.Out, srcReader)
         return err
     }
 
@@ -56,7 +59,7 @@ func (self *Drive) DownloadRevision(args DownloadRevisionArgs) (err error) {
     defer outFile.Close()
 
     // Save file to disk
-    bytes, err := io.Copy(outFile, res.Body)
+    bytes, err := io.Copy(outFile, srcReader)
     if err != nil {
         return fmt.Errorf("Failed saving file: %s", err)
     }
diff --git a/drive/update.go b/drive/update.go
index 4fa33ed..806b705 100644
--- a/drive/update.go
+++ b/drive/update.go
@@ -6,12 +6,13 @@ import (
     "os"
     "io"
     "path/filepath"
+    "google.golang.org/api/googleapi"
     "google.golang.org/api/drive/v3"
-    "golang.org/x/net/context"
 )
 
 type UpdateArgs struct {
     Out io.Writer
+    Progress io.Writer
     Id string
     Path string
     Name string
@@ -20,7 +21,7 @@ type UpdateArgs struct {
     Recursive bool
     Stdin bool
     Share bool
-    NoProgress bool
+    ChunkSize int64
 }
 
 func (self *Drive) Update(args UpdateArgs) (err error) {
@@ -58,7 +59,13 @@ func (self *Drive) Update(args UpdateArgs) (err error) {
     // Set parent folders
     dstFile.Parents = args.Parents
 
-    f, err := self.service.Files.Update(args.Id, dstFile).ResumableMedia(context.Background(), srcFile, srcFileInfo.Size(), dstFile.MimeType).Do()
+    // Chunk size option
+    chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
+
+    // Wrap file in progress reader
+    srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+
+    f, err := self.service.Files.Update(args.Id, dstFile).Media(srcReader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }
diff --git a/drive/upload.go b/drive/upload.go
index 03cb5ad..0c1df0c 100644
--- a/drive/upload.go
+++ b/drive/upload.go
@@ -12,13 +12,13 @@ import (
 
 type UploadFileArgs struct {
     Out io.Writer
+    Progress io.Writer
     Path string
     Name string
     Parents []string
     Mime string
     Recursive bool
     Share bool
-    NoProgress bool
     ChunkSize int64
 }
 
@@ -60,7 +60,10 @@ func (self *Drive) Upload(args UploadFileArgs) (err error) {
     // Chunk size option
     chunkSize := googleapi.ChunkSize(int(args.ChunkSize))
 
-    f, err := self.service.Files.Create(dstFile).Media(srcFile, chunkSize).Do()
+    // Wrap file in progress reader
+    srcReader := getProgressReader(srcFile, args.Progress, srcFileInfo.Size())
+
+    f, err := self.service.Files.Create(dstFile).Media(srcReader, chunkSize).Do()
     if err != nil {
         return fmt.Errorf("Failed to upload file: %s", err)
     }
diff --git a/gdrive.go b/gdrive.go
index 8754d90..8698911 100644
--- a/gdrive.go
+++ b/gdrive.go
@@ -274,6 +274,12 @@ func main() {
                         Description: "Share file",
                         OmitValue: true,
                     },
+                    cli.IntFlag{
+                        Name: "chunksize",
+                        Patterns: []string{"--chunksize"},
+                        Description: fmt.Sprintf("Set chunk size in bytes, default: %d", DefaultUploadChunkSize),
+                        DefaultValue: DefaultUploadChunkSize,
+                    },
                 },
             },
         },
diff --git a/handlers_drive.go b/handlers_drive.go
index 8245481..00ff781 100644
--- a/handlers_drive.go
+++ b/handlers_drive.go
@@ -3,6 +3,8 @@ package main
 import (
 	"fmt"
 	"os"
+	"io"
+	"io/ioutil"
     "./cli"
 	"./auth"
 	"./drive"
@@ -33,7 +35,7 @@ func downloadHandler(ctx cli.Context) {
         Id: args.String("id"),
         Force: args.Bool("force"),
         Stdout: args.Bool("stdout"),
-        NoProgress: args.Bool("noProgress"),
+        Progress: progressWriter(args.Bool("noProgress")),
     })
     checkErr(err)
 }
@@ -46,7 +48,7 @@ func downloadRevisionHandler(ctx cli.Context) {
         RevisionId: args.String("revisionId"),
         Force: args.Bool("force"),
         Stdout: args.Bool("stdout"),
-        NoProgress: args.Bool("noProgress"),
+        Progress: progressWriter(args.Bool("noProgress")),
     })
     checkErr(err)
 }
@@ -55,13 +57,13 @@ func uploadHandler(ctx cli.Context) {
     args := ctx.Args()
     err := newDrive(args).Upload(drive.UploadFileArgs{
         Out: os.Stdout,
+        Progress: progressWriter(args.Bool("noProgress")),
         Path: args.String("path"),
         Name: args.String("name"),
         Parents: args.StringSlice("parent"),
         Mime: args.String("mime"),
         Recursive: args.Bool("recursive"),
         Share: args.Bool("share"),
-        NoProgress: args.Bool("noProgress"),
         ChunkSize: args.Int64("chunksize"),
     })
     checkErr(err)
@@ -92,7 +94,8 @@ func updateHandler(ctx cli.Context) {
         Mime: args.String("mime"),
         Stdin: args.Bool("stdin"),
         Share: args.Bool("share"),
-        NoProgress: args.Bool("noProgress"),
+        Progress: progressWriter(args.Bool("noProgress")),
+        ChunkSize: args.Int64("chunksize"),
     })
     checkErr(err)
 }
@@ -230,3 +233,10 @@ func authCodePrompt(url string) func() string {
         return code
     }
 }
+
+func progressWriter(discard bool) io.Writer {
+    if discard {
+        return ioutil.Discard
+    }
+    return os.Stderr
+}