drive/download.go

package drive

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"google.golang.org/api/drive/v3"
	"google.golang.org/api/googleapi"
)

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
}

func (self *Drive) Download(args DownloadArgs) error {
	if args.Recursive {
		return self.downloadRecursive(args)
	}

	f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
	if err != nil {
		return fmt.Errorf("Failed to get file: %s", err)
	}

	if isDir(f) {
		return fmt.Errorf("'%s' is a directory, use --recursive to download directories", f.Name)
	}

	if !isBinary(f) {
		return fmt.Errorf("'%s' is a google document and must be exported, see the export command", f.Name)
	}

	bytes, rate, err := self.downloadBinary(f, args)
	if err != nil {
		return err
	}

	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
}

type DownloadQueryArgs struct {
	Out       io.Writer
	Progress  io.Writer
	Query     string
	Path      string
	Force     bool
	Skip      bool
	Recursive bool
}

func (self *Drive) DownloadQuery(args DownloadQueryArgs) error {
	listArgs := listAllFilesArgs{
		query:  args.Query,
		fields: []googleapi.Field{"nextPageToken", "files(id,name,mimeType,size,md5Checksum)"},
	}
	files, err := self.listAllFiles(listArgs)
	if err != nil {
		return fmt.Errorf("Failed to list files: %s", err)
	}

	downloadArgs := DownloadArgs{
		Out:      args.Out,
		Progress: args.Progress,
		Path:     args.Path,
		Force:    args.Force,
		Skip:     args.Skip,
	}

	for _, f := range files {
		if isDir(f) && args.Recursive {
			err = self.downloadDirectory(f, downloadArgs)
		} else if isBinary(f) {
			_, _, err = self.downloadBinary(f, downloadArgs)
		}

		if err != nil {
			return err
		}
	}

	return nil
}

func (self *Drive) downloadRecursive(args DownloadArgs) error {
	f, err := self.service.Files.Get(args.Id).Fields("id", "name", "size", "mimeType", "md5Checksum").Do()
	if err != nil {
		return fmt.Errorf("Failed to get file: %s", err)
	}

	if isDir(f) {
		return self.downloadDirectory(f, args)
	} else if isBinary(f) {
		_, _, err = self.downloadBinary(f, args)
		return err
	}

	return nil
}

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()
	if err != nil {
		if isTimeoutError(err) {
			return 0, 0, fmt.Errorf("Failed to download file: timeout, no data was transferred for %v", args.Timeout)
		}
		return 0, 0, fmt.Errorf("Failed to download file: %s", err)
	}

	// Close body on function exit
	defer res.Body.Close()

	// Path to file
	fpath := filepath.Join(args.Path, f.Name)

	if !args.Stdout {
		fmt.Fprintf(args.Out, "Downloading %s -> %s\n", f.Name, fpath)
	}

	return self.saveFile(saveFileArgs{
		out:           args.Out,
		body:          timeoutReaderWrapper(res.Body),
		contentLength: res.ContentLength,
		fpath:         fpath,
		force:         args.Force,
		skip:          args.Skip,
		stdout:        args.Stdout,
		progress:      args.Progress,
	})
}

type saveFileArgs struct {
	out           io.Writer
	body          io.Reader
	contentLength int64
	fpath         string
	force         bool
	skip          bool
	stdout        bool
	progress      io.Writer
}

func (self *Drive) saveFile(args saveFileArgs) (int64, int64, error) {
	// Wrap response body in progress reader
	srcReader := getProgressReader(args.body, args.progress, args.contentLength)

	if args.stdout {
		// Write file content to stdout
		_, err := io.Copy(args.out, srcReader)
		return 0, 0, err
	}

	// Check if file exists to force
	if !args.skip && !args.force && fileExists(args.fpath) {
		return 0, 0, fmt.Errorf("File '%s' already exists, use --force to overwrite or --skip to skip", args.fpath)
	}

	//Check if file exists to skip
	if args.skip && fileExists(args.fpath) {
		fmt.Printf("File '%s' already exists, skipping\n", args.fpath)
		return 0, 0, nil
	}

	// Ensure any parent directories exists
	if err := mkdir(args.fpath); err != nil {
		return 0, 0, err
	}

	// Download to tmp file
	tmpPath := args.fpath + ".incomplete"

	// Create new file
	outFile, err := os.Create(tmpPath)
	if err != nil {
		return 0, 0, fmt.Errorf("Unable to create new file: %s", err)
	}

	started := time.Now()

	// Save file to disk
	bytes, err := io.Copy(outFile, srcReader)
	if err != nil {
		outFile.Close()
		os.Remove(tmpPath)
		return 0, 0, fmt.Errorf("Failed saving file: %s", err)
	}

	// Calculate average download rate
	rate := calcRate(bytes, started, time.Now())

	// Close File
	outFile.Close()

	// Rename tmp file to proper filename
	return bytes, rate, os.Rename(tmpPath, args.fpath)
}

func (self *Drive) downloadDirectory(parent *drive.File, args DownloadArgs) error {
	listArgs := listAllFilesArgs{
		query:  fmt.Sprintf("'%s' in parents", parent.Id),
		fields: []googleapi.Field{"nextPageToken", "files(id,name)"},
	}
	files, err := self.listAllFiles(listArgs)
	if err != nil {
		return fmt.Errorf("Failed listing files: %s", err)
	}

	newPath := filepath.Join(args.Path, parent.Name)

	for _, f := range files {
		// Copy args and update changed fields
		newArgs := args
		newArgs.Path = newPath
		newArgs.Id = f.Id
		newArgs.Stdout = false

		err = self.downloadRecursive(newArgs)
		if err != nil {
			return err
		}
	}

	return nil
}

func isDir(f *drive.File) bool {
	return f.MimeType == DirectoryMimeType
}

func isBinary(f *drive.File) bool {
	return f.Md5Checksum != ""
}