download/google_drive.go

120 lines
2.3 KiB
Go

package download
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"os"
"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
}
func (g GoogleDrive) DownloadByFileID(id, name, location string) (err error) {
fmt.Printf("[google drive] start download file %s \n", name)
defer fmt.Printf("[google drive] done download file %s \n", name)
url := fmt.Sprintf("https://drive.google.com/uc?export=download&id=%s", id)
// Create the file
out, err := os.Create(location)
if err != nil {
return err
}
defer out.Close()
// 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)
}
// Writer the body to file
_, err = io.Copy(out, resp.Body)
if err != nil {
return err
}
return nil
}