Add delete option for upload and downloads

Foxsky28 · Feb 16, 2016 · 5d56138 · 5d56138
1 parent a44ec78
commit 5d56138
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 0 deletions.
diff --git a/drive/delete.go b/drive/delete.go
@@ -29,3 +29,11 @@ func (self *Drive) Delete(args DeleteArgs) error {
     fmt.Fprintf(args.Out, "Deleted '%s'\n", f.Name)
     return nil
 }
+
+func (self *Drive) deleteFile(fileId string) error {
+    err := self.service.Files.Delete(fileId).Do()
+    if err != nil {
+        return fmt.Errorf("Failed to delete file: %s", err)
+    }
+    return nil
+}
diff --git a/drive/download.go b/drive/download.go
@@ -17,6 +17,7 @@ type DownloadArgs struct {
     Path string
     Force bool
     Recursive bool
+    Delete bool
     Stdout bool
 }
 
@@ -43,6 +44,17 @@ func (self *Drive) Download(args DownloadArgs) error {
     if !args.Stdout {
         fmt.Fprintf(args.Out, "Downloaded %s at %s/s, total %s\n", f.Id, formatSize(rate, false), formatSize(bytes, false))
     }
+
+    if args.Delete {
+        err = self.deleteFile(args.Id)
+        if err != nil {
+            return fmt.Errorf("Failed to delete file: %s", err)
+        }
+
+        if !args.Stdout {
+            fmt.Fprintf(args.Out, "Removed %s\n", args.Id)
+        }
+    }
     return err
 }
 

diff --git a/drive/upload.go b/drive/upload.go
@@ -20,6 +20,7 @@ type UploadArgs struct {
     Mime string
     Recursive bool
     Share bool
+    Delete bool
     ChunkSize int64
 }
 
@@ -52,6 +53,15 @@ func (self *Drive) Upload(args UploadArgs) error {
 
         fmt.Fprintf(args.Out, "File is readable by anyone at %s\n", f.WebContentLink)
     }
+
+    if args.Delete {
+        err = os.Remove(args.Path)
+        if err != nil {
+            return fmt.Errorf("Failed to delete file: %s", err)
+        }
+        fmt.Fprintf(args.Out, "Removed %s\n", args.Path)
+    }
+
     return nil
 }
 

diff --git a/gdrive.go b/gdrive.go
@@ -100,6 +100,12 @@ func main() {
                         Patterns: []string{"--path"},
                         Description: "Download path",
                     },
+                    cli.BoolFlag{
+                        Name: "delete",
+                        Patterns: []string{"--delete"},
+                        Description: "Delete remote file when download is successful",
+                        OmitValue: true,
+                    },
                     cli.BoolFlag{
                         Name: "noProgress",
                         Patterns: []string{"--no-progress"},
@@ -155,6 +161,12 @@ func main() {
                         Description: "Share file",
                         OmitValue: true,
                     },
+                    cli.BoolFlag{
+                        Name: "delete",
+                        Patterns: []string{"--delete"},
+                        Description: "Delete local file when upload is successful",
+                        OmitValue: true,
+                    },
                     cli.IntFlag{
                         Name: "chunksize",
                         Patterns: []string{"--chunksize"},

diff --git a/handlers_drive.go b/handlers_drive.go
@@ -46,11 +46,13 @@ func listChangesHandler(ctx cli.Context) {
 
 func downloadHandler(ctx cli.Context) {
     args := ctx.Args()
+    checkDownloadArgs(args)
     err := newDrive(args).Download(drive.DownloadArgs{
         Out: os.Stdout,
         Id: args.String("fileId"),
         Force: args.Bool("force"),
         Path: args.String("path"),
+        Delete: args.Bool("delete"),
         Recursive: args.Bool("recursive"),
         Stdout: args.Bool("stdout"),
         Progress: progressWriter(args.Bool("noProgress")),
@@ -90,6 +92,7 @@ func downloadRevisionHandler(ctx cli.Context) {
 
 func uploadHandler(ctx cli.Context) {
     args := ctx.Args()
+    checkUploadArgs(args)
     err := newDrive(args).Upload(drive.UploadArgs{
         Out: os.Stdout,
         Progress: progressWriter(args.Bool("noProgress")),
@@ -99,6 +102,7 @@ func uploadHandler(ctx cli.Context) {
         Mime: args.String("mime"),
         Recursive: args.Bool("recursive"),
         Share: args.Bool("share"),
+        Delete: args.Bool("delete"),
         ChunkSize: args.Int64("chunksize"),
     })
     checkErr(err)
@@ -366,3 +370,19 @@ func conflictResolution(args cli.Arguments) drive.ConflictResolution {
 
     return drive.NoResolution
 }
+
+func checkUploadArgs(args cli.Arguments) {
+    if args.Bool("recursive") && args.Bool("delete") {
+        ExitF("--delete is not allowed for recursive uploads")
+    }
+
+    if args.Bool("recursive") && args.Bool("share") {
+        ExitF("--share is not allowed for recursive uploads")
+    }
+}
+
+func checkDownloadArgs(args cli.Arguments) {
+    if args.Bool("recursive") && args.Bool("delete") {
+        ExitF("--delete is not allowed for recursive downloads")
+    }
+}