137 lines
2.8 KiB
Go
137 lines
2.8 KiB
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) (filename 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 filename, nil
|
|
}
|