Update README.md

README.md

@@ -1,2 +1,4 @@
 # download
 
+
+

google_drive.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"strings"
 
 	"github.com/tanaikech/go-getfilelist"
 	"google.golang.org/api/drive/v3"
@@ -45,18 +46,23 @@ func NewGoogleDrive(apiKey string) (*GoogleDrive, error) {
 	return &ggdrive, nil
 }
 
-func (g GoogleDrive) GetFilesByFolderID(folderID string) ([]GoogleDriveFile, error) {
+func (g GoogleDrive) GetFilesByFolderID(folderID string, mimeType []string) ([]GoogleDriveFile, error) {
 	var (
 		result = make([]GoogleDriveFile, 0)
 	)
 
-	res, err := getfilelist.Folder(folderID).Fields("files(name,id)").Do(g.svc)
+	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 subfolder
+	// 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
@@ -79,36 +85,52 @@ func (g GoogleDrive) GetFilesByFolderID(folderID string) ([]GoogleDriveFile, err
 	return result, nil
 }
 
-func (g GoogleDrive) DownloadByFileID(id, name, location string) (err error) {
+// DownloadByFileID download file into specific path
-	fmt.Printf("[google drive] start download file %s \n", name)
+func (g GoogleDrive) DownloadByFileID(id, path string) (filename string, err error) {
-	defer fmt.Printf("[google drive] done download file %s \n", name)
+	// 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)
 
-	// 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
+		return "", err
 	}
 	defer resp.Body.Close()
 
 	// Check server response
 	if resp.StatusCode != http.StatusOK {
-		return fmt.Errorf("bad status: %s", resp.Status)
+		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 "", err
 	}
 
-	return nil
+	return filename, nil
 }

util.go

@@ -0,0 +1,13 @@
+package download
+
+import (
+	"mime"
+)
+
+func getFileNameFromHeaderContentDisposition(value string) string {
+	disposition, params, err := mime.ParseMediaType(value)
+	if err != nil || disposition != "attachment" {
+		return ""
+	}
+	return params["filename"]
+}