feat(google_drive): chunk upload (close #2241)

pull/2258/head
Noah Hsu 2022-11-07 20:58:52 +08:00
parent 8241f0999a
commit 6639cab1ae
3 changed files with 37 additions and 3 deletions

View File

@ -36,6 +36,9 @@ func (d *GoogleDrive) Init(ctx context.Context, storage model.Storage) error {
if err != nil {
return err
}
if d.ChunkSize == 0 {
d.ChunkSize = 5
}
return d.refreshToken()
}
@ -160,9 +163,13 @@ func (d *GoogleDrive) Put(ctx context.Context, dstDir model.Obj, stream model.Fi
return fmt.Errorf("%s: %v", e.Error.Message, e.Error.Errors)
}
putUrl := res.Header().Get("location")
_, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) {
req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser())
}, nil)
if stream.GetSize() < d.ChunkSize*1024*1024 {
_, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) {
req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser())
}, nil)
} else {
err = d.chunkUpload(ctx, stream, putUrl)
}
return err
}

View File

@ -12,6 +12,7 @@ type Addition struct {
OrderDirection string `json:"order_direction" type:"select" options:"asc,desc"`
ClientID string `json:"client_id" required:"true" default:"202264815644.apps.googleusercontent.com"`
ClientSecret string `json:"client_secret" required:"true" default:"X4Z3ca8xfWDb1Voo-F9a7ZxJ"`
ChunkSize int64 `json:"chunk_size" default:"5" help:"chunk size while uploading (unit: MB)"`
}
var config = driver.Config{

View File

@ -1,10 +1,14 @@
package google_drive
import (
"context"
"fmt"
"io"
"net/http"
"strconv"
"github.com/alist-org/alist/v3/drivers/base"
"github.com/alist-org/alist/v3/internal/model"
"github.com/go-resty/resty/v2"
log "github.com/sirupsen/logrus"
)
@ -95,3 +99,25 @@ func (d *GoogleDrive) getFiles(id string) ([]File, error) {
}
return res, nil
}
func (d *GoogleDrive) chunkUpload(ctx context.Context, stream model.FileStreamer, url string) error {
var defaultChunkSize = d.ChunkSize * 1024 * 1024
var finish int64 = 0
for finish < stream.GetSize() {
chunkSize := stream.GetSize() - finish
if chunkSize > defaultChunkSize {
chunkSize = defaultChunkSize
}
_, err := d.request(url, http.MethodPut, func(req *resty.Request) {
req.SetHeaders(map[string]string{
"Content-Length": strconv.FormatInt(chunkSize, 10),
"Content-Range": fmt.Sprintf("bytes %d-%d/%d", finish, finish+chunkSize-1, stream.GetSize()),
}).SetBody(io.LimitReader(stream.GetReadCloser(), chunkSize))
}, nil)
if err != nil {
return err
}
finish += chunkSize
}
return nil
}