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