diff --git a/google_drive.go b/google_drive.go index 28f19c6..7c2c417 100644 --- a/google_drive.go +++ b/google_drive.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "os" + "strings" "github.com/tanaikech/go-getfilelist" "google.golang.org/api/drive/v3" @@ -84,19 +85,14 @@ func (g GoogleDrive) GetFilesByFolderID(folderID string, mimeType []string) ([]G 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) +// 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) - // 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 { @@ -109,6 +105,27 @@ func (g GoogleDrive) DownloadByFileID(id, name, location string) (err error) { 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 { diff --git a/util.go b/util.go new file mode 100644 index 0000000..55c89ac --- /dev/null +++ b/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"] +}