support teamdrives

drive/delete.go

@@ -12,7 +12,7 @@ type DeleteArgs struct {
 }
 
 func (self *Drive) Delete(args DeleteArgs) error {
-	f, err := self.service.Files.Get(args.Id).Fields("name", "mimeType").Do()
+	f, err := self.service.Files.Get(args.Id).SupportsAllDrives(true).Fields("name", "mimeType").Do()
 	if err != nil {
 		return fmt.Errorf("Failed to get file: %s", err)
 	}
@@ -21,7 +21,7 @@ func (self *Drive) Delete(args DeleteArgs) error {
 		return fmt.Errorf("'%s' is a directory, use the 'recursive' flag to delete directories", f.Name)
 	}
 
-	err = self.service.Files.Delete(args.Id).Do()
+	err = self.service.Files.Delete(args.Id).SupportsAllDrives(true).Do()
 	if err != nil {
 		return fmt.Errorf("Failed to delete file: %s", err)
 	}
@@ -31,7 +31,7 @@ func (self *Drive) Delete(args DeleteArgs) error {
 }
 
 func (self *Drive) deleteFile(fileId string) error {
-	err := self.service.Files.Delete(fileId).Do()
+	err := self.service.Files.Delete(fileId).SupportsAllDrives(true).Do()
 	if err != nil {
 		return fmt.Errorf("Failed to delete file: %s", err)
 	}

drive/download.go

@@ -3,6 +3,7 @@ package drive
 import (
 	"fmt"
 	"io"
+	"math"
 	"os"
 	"path/filepath"
 	"time"
@@ -12,26 +13,32 @@ import (
 )
 
 type DownloadArgs struct {
-	Out       io.Writer
-	Progress  io.Writer
-	Id        string
-	Path      string
-	Force     bool
-	Skip      bool
-	Recursive bool
-	Delete    bool
-	Stdout    bool
-	Timeout   time.Duration
+	Out                 io.Writer
+	Progress            io.Writer
+	Id                  string
+	Path                string
+	Force               bool
+	Skip                bool
+	Recursive           bool
+	Delete              bool
+	Stdout              bool
+	Timeout             time.Duration
+	RecursiveExtraQuery string // add extra query
+	IsAsyncDownload     bool
+	LimitPerSec         int64
 }
 
 func (self *Drive) Download(args DownloadArgs) error {
 	if args.Recursive {
-		return self.downloadRecursive(args)
+		self.ResetDownloadTime()
+		self.downloadRecursive(args)
+		self.waitGroup.Wait()
+		return self.downloadErr
 	}
 
-	f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
+	f, err := self.service.Files.Get(args.Id).SupportsAllDrives(true).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
 	if err != nil {
-		return fmt.Errorf("Failed to get file: %s", err)
+		return fmt.Errorf("Failed to get file: %s, err:", args.Id, err)
 	}
 
 	if isDir(f) {
@@ -65,13 +72,16 @@ func (self *Drive) Download(args DownloadArgs) error {
 }
 
 type DownloadQueryArgs struct {
-	Out       io.Writer
-	Progress  io.Writer
-	Query     string
-	Path      string
-	Force     bool
-	Skip      bool
-	Recursive bool
+	Out                 io.Writer
+	Progress            io.Writer
+	Query               string
+	Path                string
+	Force               bool
+	Skip                bool
+	Recursive           bool
+	RecursiveExtraQuery string
+	IsAsyncDownload     bool
+	LimitPerSec         int64
 }
 
 func (self *Drive) DownloadQuery(args DownloadQueryArgs) error {
@@ -85,13 +95,18 @@ func (self *Drive) DownloadQuery(args DownloadQueryArgs) error {
 	}
 
 	downloadArgs := DownloadArgs{
-		Out:      args.Out,
-		Progress: args.Progress,
-		Path:     args.Path,
-		Force:    args.Force,
-		Skip:     args.Skip,
+		Out:                 args.Out,
+		Progress:            args.Progress,
+		Path:                args.Path,
+		Force:               args.Force,
+		Skip:                args.Skip,
+		RecursiveExtraQuery: args.RecursiveExtraQuery,
+		IsAsyncDownload:     args.IsAsyncDownload,
+		LimitPerSec:         args.LimitPerSec,
 	}
 
+	self.ResetDownloadTime()
+
 	for _, f := range files {
 		if isDir(f) && args.Recursive {
 			err = self.downloadDirectory(f, downloadArgs)
@@ -104,30 +119,64 @@ func (self *Drive) DownloadQuery(args DownloadQueryArgs) error {
 		}
 	}
 
-	return nil
+	self.waitGroup.Wait()
+
+	return self.downloadErr
 }
 
 func (self *Drive) downloadRecursive(args DownloadArgs) error {
-	f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
+	self.downlaodCount++
+	perid := time.Now().Unix() - self.downloadStartUnix
+	if perid < 1 {
+		perid = 1
+	}
+	limit := perid * args.LimitPerSec
+	if limit < 0 {
+		limit = math.MaxInt64
+	}
+	for self.downlaodCount > limit {
+		gap := self.downlaodCount - limit
+		time.Sleep(time.Duration(gap) * time.Second)
+		perid = time.Now().Unix() - self.downloadStartUnix
+		limit = perid * args.LimitPerSec
+	}
+
+	f, err := self.service.Files.Get(args.Id).SupportsAllDrives(true).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
 	if err != nil {
-		return fmt.Errorf("Failed to get file: %s", err)
+		return fmt.Errorf("Failed to get file: %s, err:", args.Id, err)
+	}
+
+	if args.IsAsyncDownload {
+		self.waitGroup.Add(1)
+		go func() {
+			self.doDownloadRecursive(f, args)
+			self.waitGroup.Done()
+		}()
+	} else {
+		self.doDownloadRecursive(f, args)
 	}
 
+	return nil
+}
+
+func (self *Drive) doDownloadRecursive(f *drive.File, args DownloadArgs) {
+	var err error
 	if isDir(f) {
-		return self.downloadDirectory(f, args)
+		err = self.downloadDirectory(f, args)
 	} else if isBinary(f) {
 		_, _, err = self.downloadBinary(f, args)
-		return err
 	}
 
-	return nil
+	if err != nil {
+		fmt.Errorf("%s \n %s", self.downloadErr, err.Error())
+		//fmt.Println("Failed to download:", err)
+	}
 }
 
 func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) (int64, int64, error) {
 	// Get timeout reader wrapper and context
 	timeoutReaderWrapper, ctx := getTimeoutReaderWrapperContext(args.Timeout)
-
-	res, err := self.service.Files.Get(f.Id).Context(ctx).Download()
+	res, err := self.service.Files.Get(f.Id).SupportsAllDrives(true).Context(ctx).Download()
 	if err != nil {
 		if isTimeoutError(err) {
 			return 0, 0, fmt.Errorf("Failed to download file: timeout, no data was transferred for %v", args.Timeout)
@@ -140,6 +189,7 @@ func (self *Drive) downloadBinary(f *drive.File, args DownloadArgs) (int64, int6
 
 	// Path to file
 	fpath := filepath.Join(args.Path, f.Name)
+	fmt.Println("fpath:", fpath)
 
 	if !args.Stdout {
 		fmt.Fprintf(args.Out, "Downloading %s -> %s\n", f.Name, fpath)
@@ -197,6 +247,8 @@ func (self *Drive) saveFile(args saveFileArgs) (int64, int64, error) {
 	// Download to tmp file
 	tmpPath := args.fpath + ".incomplete"
 
+	fmt.Println("tmpPath:", tmpPath)
+
 	// Create new file
 	outFile, err := os.Create(tmpPath)
 	if err != nil {
@@ -225,7 +277,7 @@ func (self *Drive) saveFile(args saveFileArgs) (int64, int64, error) {
 
 func (self *Drive) downloadDirectory(parent *drive.File, args DownloadArgs) error {
 	listArgs := listAllFilesArgs{
-		query:  fmt.Sprintf("'%s' in parents", parent.Id),
+		query:  fmt.Sprintf("'%s' in parents %s", parent.Id, args.RecursiveExtraQuery),
 		fields: []googleapi.Field{"nextPageToken", "files(id,name)"},
 	}
 	files, err := self.listAllFiles(listArgs)

drive/drive.go

@@ -1,12 +1,19 @@
 package drive
 
 import (
-	"google.golang.org/api/drive/v3"
 	"net/http"
+	"sync"
+	"time"
+
+	"google.golang.org/api/drive/v3"
 )
 
 type Drive struct {
-	service *drive.Service
+	service           *drive.Service
+	downloadStartUnix int64
+	downlaodCount     int64
+	waitGroup         sync.WaitGroup
+	downloadErr       error
 }
 
 func New(client *http.Client) (*Drive, error) {
@@ -15,5 +22,10 @@ func New(client *http.Client) (*Drive, error) {
 		return nil, err
 	}
 
-	return &Drive{service}, nil
+	return &Drive{service, 0, 0, sync.WaitGroup{}, nil}, nil
+}
+
+func (d *Drive) ResetDownloadTime() {
+	d.downloadStartUnix = time.Now().Unix()
+	d.downlaodCount = 0
 }

drive/export.go

@@ -25,7 +25,7 @@ type ExportArgs struct {
 }
 
 func (self *Drive) Export(args ExportArgs) error {
-	f, err := self.service.Files.Get(args.Id).Fields("name", "mimeType").Do()
+	f, err := self.service.Files.Get(args.Id).SupportsTeamDrives(true).Fields("name", "mimeType").Do()
 	if err != nil {
 		return fmt.Errorf("Failed to get file: %s", err)
 	}

drive/info.go

@@ -2,8 +2,9 @@ package drive
 
 import (
 	"fmt"
-	"google.golang.org/api/drive/v3"
 	"io"
+
+	"google.golang.org/api/drive/v3"
 )
 
 type FileInfoArgs struct {
@@ -13,7 +14,7 @@ type FileInfoArgs struct {
 }
 
 func (self *Drive) Info(args FileInfoArgs) error {
-	f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "createdTime", "modifiedTime", "md5Checksum", "mimeType", "parents", "shared", "description", "webContentLink", "webViewLink").Do()
+	f, err := self.service.Files.Get(args.Id).SupportsAllDrives(true).Fields("id", "name", "size", "createdTime", "modifiedTime", "md5Checksum", "mimeType", "parents", "shared", "description", "webContentLink", "webViewLink").Do()
 	if err != nil {
 		return fmt.Errorf("Failed to get file: %s", err)
 	}

drive/list.go

@@ -2,11 +2,12 @@ package drive
 
 import (
 	"fmt"
+	"io"
+	"text/tabwriter"
+
 	"golang.org/x/net/context"
 	"google.golang.org/api/drive/v3"
 	"google.golang.org/api/googleapi"
-	"io"
-	"text/tabwriter"
 )
 
 type ListFilesArgs struct {
@@ -73,8 +74,8 @@ func (self *Drive) listAllFiles(args listAllFilesArgs) ([]*drive.File, error) {
 	}
 
 	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 {
+	fmt.Println(args.query)
+	err := self.service.Files.List().Corpora("allDrives").IncludeItemsFromAllDrives(true).SupportsAllDrives(true).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

drive/mkdir.go

@@ -2,8 +2,9 @@ package drive
 
 import (
 	"fmt"
-	"google.golang.org/api/drive/v3"
 	"io"
+
+	"google.golang.org/api/drive/v3"
 )
 
 const DirectoryMimeType = "application/vnd.google-apps.folder"
@@ -35,7 +36,7 @@ func (self *Drive) mkdir(args MkdirArgs) (*drive.File, error) {
 	dstFile.Parents = args.Parents
 
 	// Create directory
-	f, err := self.service.Files.Create(dstFile).Do()
+	f, err := self.service.Files.Create(dstFile).SupportsAllDrives(true).Do()
 	if err != nil {
 		return nil, fmt.Errorf("Failed to create directory: %s", err)
 	}

drive/sync.go

@@ -2,16 +2,17 @@ package drive
 
 import (
 	"fmt"
-	"github.com/sabhiram/go-git-ignore"
-	"github.com/soniakeys/graph"
-	"google.golang.org/api/drive/v3"
-	"google.golang.org/api/googleapi"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
 	"text/tabwriter"
 	"time"
+
+	"github.com/sabhiram/go-git-ignore"
+	"github.com/soniakeys/graph"
+	"google.golang.org/api/drive/v3"
+	"google.golang.org/api/googleapi"
 )
 
 const DefaultIgnoreFile = ".gdriveignore"
@@ -86,7 +87,7 @@ func (self *Drive) prepareSyncFiles(localPath string, root *drive.File, cmp File
 }
 
 func (self *Drive) isSyncFile(id string) (bool, error) {
-	f, err := self.service.Files.Get(id).Fields("appProperties").Do()
+	f, err := self.service.Files.Get(id).SupportsAllDrives(true).Fields("appProperties").Do()
 	if err != nil {
 		return false, fmt.Errorf("Failed to get file: %s", err)
 	}