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 { if err != nil {
return err return err
} }
if d.ChunkSize == 0 {
d.ChunkSize = 5
}
return d.refreshToken() 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) return fmt.Errorf("%s: %v", e.Error.Message, e.Error.Errors)
} }
putUrl := res.Header().Get("location") putUrl := res.Header().Get("location")
if stream.GetSize() < d.ChunkSize*1024*1024 {
_, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) { _, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) {
req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser()) req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser())
}, nil) }, nil)
} else {
err = d.chunkUpload(ctx, stream, putUrl)
}
return err return err
} }

View File

@ -12,6 +12,7 @@ type Addition struct {
OrderDirection string `json:"order_direction" type:"select" options:"asc,desc"` OrderDirection string `json:"order_direction" type:"select" options:"asc,desc"`
ClientID string `json:"client_id" required:"true" default:"202264815644.apps.googleusercontent.com"` ClientID string `json:"client_id" required:"true" default:"202264815644.apps.googleusercontent.com"`
ClientSecret string `json:"client_secret" required:"true" default:"X4Z3ca8xfWDb1Voo-F9a7ZxJ"` 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{ var config = driver.Config{

View File

@ -1,10 +1,14 @@
package google_drive package google_drive
import ( import (
"context"
"fmt" "fmt"
"io"
"net/http" "net/http"
"strconv"
"github.com/alist-org/alist/v3/drivers/base" "github.com/alist-org/alist/v3/drivers/base"
"github.com/alist-org/alist/v3/internal/model"
"github.com/go-resty/resty/v2" "github.com/go-resty/resty/v2"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@ -95,3 +99,25 @@ func (d *GoogleDrive) getFiles(id string) ([]File, error) {
} }
return res, nil 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
}