download/google_drive.go

package download

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"

	"github.com/tanaikech/go-getfilelist"
	"google.golang.org/api/drive/v3"
	"google.golang.org/api/option"
)

type GoogleDrive struct {
	apiKey string
	svc    *drive.Service
}

type GoogleDriveFile struct {
	ID   string
	Name string
}

var ggdrive GoogleDrive

func NewGoogleDrive(apiKey string) (*GoogleDrive, error) {
	if apiKey == "" {
		err := errors.New("missing google api key")
		return nil, err
	}

	// init google drive
	svc, err := drive.NewService(context.Background(), option.WithAPIKey(apiKey))
	if err != nil {
		err = fmt.Errorf("unable to retrieve Drive: %v", err)
		return nil, err
	}

	// assign
	ggdrive.apiKey = apiKey
	ggdrive.svc = svc

	return &ggdrive, nil
}

func (g GoogleDrive) GetFilesByFolderID(folderID string, mimeType []string) ([]GoogleDriveFile, error) {
	var (
		result = make([]GoogleDriveFile, 0)
	)

	chain := getfilelist.Folder(folderID).Fields("files(id,name)")
	if len(mimeType) > 0 {
		chain.MimeType(mimeType)
	}

	res, err := chain.Do(g.svc)
	if err != nil {
		err := fmt.Errorf("error when fetch folder %s: %s", folderID, err.Error())
		return result, err
	}

	// if there is more than 1 sub folder
	if len(res.FileList) != 1 {
		err := fmt.Errorf("folder %s have more than 1 subfolder", folderID)
		return result, err
	}

	// check list files
	fol := res.FileList[0]
	if len(fol.Files) == 0 {
		return result, nil
	}

	// collect files
	for _, f := range fol.Files {
		result = append(result, GoogleDriveFile{
			ID:   f.Id,
			Name: f.Name,
		})
	}

	return result, nil
}

// DownloadByFileID download file into specific path
func (g GoogleDrive) DownloadByFileID(id, path string) (err error) {
	// trim "/" in the end of path
	path = strings.TrimSuffix(path, "/")

	// prepare url
	url := fmt.Sprintf("https://drive.google.com/uc?export=download&id=%s", id)

	// Get the data
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Check server response
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad status: %s", resp.Status)
	}

	// get file name
	filename := getFileNameFromHeaderContentDisposition(resp.Header.Get("Content-Disposition"))

	// return err if not found
	if filename == "" {
		return errors.New("file is not existed or is in private mode, please check again")
	}

	// assign location with field name
	location := fmt.Sprintf("%s/%s", path, filename)

	// Create the file
	out, err := os.Create(location)
	if err != nil {
		return err
	}
	defer out.Close()

	fmt.Printf("[google drive] start download file %s with id %s \n", filename, id)
	defer fmt.Printf("[google drive] done download file %s with id %s \n", filename, id)

	// Writer the body to file
	_, err = io.Copy(out, resp.Body)
	if err != nil {
		return err
	}

	return nil
}