diff options
Diffstat (limited to 'drive')
| -rw-r--r-- | drive/download.go | 9 | ||||
| -rw-r--r-- | drive/progress.go | 13 | ||||
| -rw-r--r-- | drive/revision_download.go | 9 | ||||
| -rw-r--r-- | drive/update.go | 13 | ||||
| -rw-r--r-- | drive/upload.go | 7 | 
5 files changed, 40 insertions, 11 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)      } | 
