Skip to content

Commit e6991d6

Browse files
darthShadowdarthShadow
committed
drive: load all clients as specified in the config file
Signed-off-by: Anagh Kumar Baranwal <6824881+darthShadow@users.noreply.github.com>
1 parent bd9be8d commit e6991d6

File tree

2 files changed

+138
-15
lines changed

2 files changed

+138
-15
lines changed

‎backend/drive/drive.go‎

Lines changed: 54 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -848,6 +848,10 @@ type Fs struct {
848848
dirResourceKeys *sync.Map // map directory ID to resource key
849849
permissionsMu *sync.Mutex // protect the below
850850
permissions map[string]*drive.Permission // map permission IDs to Permissions
851+
852+
uploadClients []*http.Client // authorized clients for uploading
853+
uploadClientMutex *sync.Mutex // protects uploadClients and uploadClientIndex
854+
uploadClientIndex int // index of the next client to use
851855
}
852856

853857
type baseObject struct {
@@ -1286,28 +1290,39 @@ func getServiceAccountClient(ctx context.Context, opt *Options, credentialsData
12861290
}
12871291

12881292
func createOAuthClient(ctx context.Context, opt *Options, name string, m configmap.Mapper) (*http.Client, error) {
1289-
var oAuthClient *http.Client
1290-
var err error
1293+
var (
1294+
err error
1295+
oAuthClient *http.Client
1296+
)
12911297

1292-
// try loading service account credentials from env variable, then from a file
1293-
if len(opt.ServiceAccountCredentials) == 0 && opt.ServiceAccountFile != "" {
1294-
loadedCreds, err := os.ReadFile(env.ShellExpand(opt.ServiceAccountFile))
1298+
if clientID, ok := m.Get(config.ConfigClientID); ok && clientID != "" {
1299+
oAuthClient, _, err = oauthutil.NewClientWithBaseClient(ctx, name, m, driveConfig, getClient(ctx, opt))
12951300
if err != nil {
1296-
return nil, fmt.Errorf("error opening service account credentials file: %w", err)
1301+
return nil, fmt.Errorf("failed to create oauth client: %w", err)
12971302
}
1298-
opt.ServiceAccountCredentials = string(loadedCreds)
1299-
}
1300-
if opt.ServiceAccountCredentials != "" {
1303+
1304+
} else if opt.ServiceAccountCredentials != "" || opt.ServiceAccountFile != "" {
1305+
// try loading service account credentials from env variable, then from a file
1306+
if opt.ServiceAccountCredentials == "" && opt.ServiceAccountFile != "" {
1307+
loadedCreds, err := os.ReadFile(env.ShellExpand(opt.ServiceAccountFile))
1308+
if err != nil {
1309+
return nil, fmt.Errorf("error opening service account credentials file: %w", err)
1310+
}
1311+
opt.ServiceAccountCredentials = string(loadedCreds)
1312+
}
1313+
13011314
oAuthClient, err = getServiceAccountClient(ctx, opt, []byte(opt.ServiceAccountCredentials))
13021315
if err != nil {
13031316
return nil, fmt.Errorf("failed to create oauth client from service account: %w", err)
13041317
}
1318+
13051319
} else if opt.EnvAuth {
13061320
scopes := driveScopes(opt.Scope)
13071321
oAuthClient, err = google.DefaultClient(ctx, scopes...)
13081322
if err != nil {
13091323
return nil, fmt.Errorf("failed to create client from environment: %w", err)
13101324
}
1325+
13111326
} else {
13121327
oAuthClient, _, err = oauthutil.NewClientWithBaseClient(ctx, name, m, driveConfig, getClient(ctx, opt))
13131328
if err != nil {
@@ -1411,8 +1426,17 @@ func newFs(ctx context.Context, name, path string, m configmap.Mapper) (*Fs, err
14111426
DirModTimeUpdatesOnWrite: false, // FIXME need to check!
14121427
}).Fill(ctx, f)
14131428

1414-
// Create a new authorized Drive client.
14151429
f.client = oAuthClient
1430+
1431+
f.uploadClients, err = loadOAuthClients(ctx, opt, name, m)
1432+
if err != nil {
1433+
fs.Errorf(f, "Failed to load upload clients: %v", err)
1434+
f.uploadClients = []*http.Client{f.client}
1435+
}
1436+
f.uploadClientIndex = 0
1437+
f.uploadClientMutex = new(sync.Mutex)
1438+
1439+
// Create a new authorized Drive client.
14161440
f.svc, err = drive.NewService(context.Background(), option.WithHTTPClient(f.client))
14171441
if err != nil {
14181442
return nil, fmt.Errorf("couldn't create Drive client: %w", err)
@@ -3232,15 +3256,21 @@ func (f *Fs) changeChunkSize(chunkSizeString string) (err error) {
32323256
}
32333257

32343258
func (f *Fs) changeServiceAccountFile(ctx context.Context, file string) (err error) {
3235-
fs.Debugf(nil, "Changing Service Account File from %s to %s", f.opt.ServiceAccountFile, file)
32363259
if file == f.opt.ServiceAccountFile {
32373260
return nil
32383261
}
32393262

3263+
fs.Debugf(nil, "Changing Service Account File from %s to %s", f.opt.ServiceAccountFile, file)
3264+
3265+
uploadClientMutex := f.uploadClientMutex
3266+
uploadClientMutex.Lock()
3267+
defer uploadClientMutex.Unlock()
3268+
32403269
oldSvc := f.svc
32413270
oldActivitySvc := f.activitySvc
32423271
oldv2Svc := f.v2Svc
32433272
oldOAuthClient := f.client
3273+
oldUploadClients := f.uploadClients
32443274
oldFile := f.opt.ServiceAccountFile
32453275
oldCredentials := f.opt.ServiceAccountCredentials
32463276

@@ -3251,6 +3281,7 @@ func (f *Fs) changeServiceAccountFile(ctx context.Context, file string) (err err
32513281
f.activitySvc = oldActivitySvc
32523282
f.v2Svc = oldv2Svc
32533283
f.client = oldOAuthClient
3284+
f.uploadClients = oldUploadClients
32543285
f.opt.ServiceAccountFile = oldFile
32553286
f.opt.ServiceAccountCredentials = oldCredentials
32563287
}
@@ -3262,15 +3293,27 @@ func (f *Fs) changeServiceAccountFile(ctx context.Context, file string) (err err
32623293
if err != nil {
32633294
return fmt.Errorf("drive: failed when making oauth client: %w", err)
32643295
}
3296+
32653297
f.client = oAuthClient
3298+
3299+
f.uploadClients, err = loadOAuthClients(ctx, &f.opt, f.name, f.m)
3300+
if err != nil {
3301+
fs.Errorf(nil, "Failed to load upload clients: %v", err)
3302+
f.uploadClients = []*http.Client{f.client}
3303+
}
3304+
f.uploadClientIndex = 0
3305+
3306+
// Create a new authorized Drive client.
32663307
f.svc, err = drive.NewService(context.Background(), option.WithHTTPClient(f.client))
32673308
if err != nil {
32683309
return fmt.Errorf("couldn't create Drive client: %w", err)
32693310
}
3311+
32703312
f.activitySvc, err = driveactivity.NewService(context.Background(), option.WithHTTPClient(f.client))
32713313
if err != nil {
32723314
return fmt.Errorf("couldn't create Drive Activity client: %w", err)
32733315
}
3316+
32743317
if f.opt.V2DownloadMinSize >= 0 {
32753318
f.v2Svc, err = drive_v2.NewService(context.Background(), option.WithHTTPClient(f.client))
32763319
if err != nil {

‎backend/drive/upload.go‎

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,17 @@ import (
1818
"io"
1919
"net/http"
2020
"net/url"
21+
"os"
2122
"strconv"
2223

2324
"github.com/rclone/rclone/fs"
25+
"github.com/rclone/rclone/fs/config"
26+
"github.com/rclone/rclone/fs/config/configmap"
2427
"github.com/rclone/rclone/fs/fserrors"
28+
"github.com/rclone/rclone/lib/env"
29+
"github.com/rclone/rclone/lib/oauthutil"
2530
"github.com/rclone/rclone/lib/readers"
31+
"golang.org/x/oauth2/google"
2632
"google.golang.org/api/drive/v3"
2733
"google.golang.org/api/googleapi"
2834
)
@@ -35,8 +41,9 @@ const (
3541
// resumableUpload is used by the generated APIs to provide resumable uploads.
3642
// It is not used by developers directly.
3743
type resumableUpload struct {
38-
f *Fs
39-
remote string
44+
f *Fs
45+
uploadClient *http.Client // HTTP client to use in making requests
46+
remote string
4047
// URI is the resumable resource destination provided by the server after specifying "&uploadType=resumable".
4148
URI string
4249
// Media is the object being uploaded.
@@ -51,6 +58,11 @@ type resumableUpload struct {
5158

5259
// Upload the io.Reader in of size bytes with contentType and info
5360
func (f *Fs) Upload(ctx context.Context, in io.Reader, size int64, contentType, fileID, remote string, info *drive.File) (*drive.File, error) {
61+
uploadClient := f.getUploadClient(ctx)
62+
if uploadClient == nil {
63+
uploadClient = f.client
64+
}
65+
5466
params := url.Values{
5567
"alt": {"json"},
5668
"uploadType": {"resumable"},
@@ -89,7 +101,7 @@ func (f *Fs) Upload(ctx context.Context, in io.Reader, size int64, contentType,
89101
if size >= 0 {
90102
req.Header.Set("X-Upload-Content-Length", fmt.Sprintf("%v", size))
91103
}
92-
res, err = f.client.Do(req)
104+
res, err = uploadClient.Do(req)
93105
if err == nil {
94106
defer googleapi.CloseBody(res)
95107
err = googleapi.CheckResponse(res)
@@ -102,6 +114,7 @@ func (f *Fs) Upload(ctx context.Context, in io.Reader, size int64, contentType,
102114
loc := res.Header.Get("Location")
103115
rx := &resumableUpload{
104116
f: f,
117+
uploadClient: uploadClient,
105118
remote: remote,
106119
URI: loc,
107120
Media: in,
@@ -132,7 +145,7 @@ func (rx *resumableUpload) makeRequest(ctx context.Context, start int64, body io
132145
func (rx *resumableUpload) transferChunk(ctx context.Context, start int64, chunk io.ReadSeeker, chunkSize int64) (int, error) {
133146
_, _ = chunk.Seek(0, io.SeekStart)
134147
req := rx.makeRequest(ctx, start, chunk, chunkSize)
135-
res, err := rx.f.client.Do(req)
148+
res, err := rx.uploadClient.Do(req)
136149
if err != nil {
137150
return 599, err
138151
}
@@ -240,3 +253,70 @@ func (rx *resumableUpload) Upload(ctx context.Context) (*drive.File, error) {
240253
}
241254
return rx.ret, nil
242255
}
256+
257+
func loadOAuthClients(ctx context.Context, opt *Options, name string, m configmap.Mapper) ([]*http.Client, error) {
258+
var (
259+
err error
260+
oAuthClient *http.Client
261+
)
262+
263+
oAuthClients := make([]*http.Client, 0, 3)
264+
265+
if clientID, ok := m.Get(config.ConfigClientID); ok && clientID != "" {
266+
oAuthClient, _, err = oauthutil.NewClientWithBaseClient(ctx, name, m, driveConfig, getClient(ctx, opt))
267+
if err != nil {
268+
return nil, fmt.Errorf("failed to create oauth client: %w", err)
269+
}
270+
oAuthClients = append(oAuthClients, oAuthClient)
271+
}
272+
273+
if opt.EnvAuth {
274+
scopes := driveScopes(opt.Scope)
275+
oAuthClient, err = google.DefaultClient(ctx, scopes...)
276+
if err != nil {
277+
return nil, fmt.Errorf("failed to create client from environment: %w", err)
278+
}
279+
oAuthClients = append(oAuthClients, oAuthClient)
280+
}
281+
282+
// try loading service account credentials from env variable, then from a file
283+
if len(opt.ServiceAccountCredentials) == 0 && opt.ServiceAccountFile != "" {
284+
loadedCreds, err := os.ReadFile(env.ShellExpand(opt.ServiceAccountFile))
285+
if err != nil {
286+
return nil, fmt.Errorf("error opening service account credentials file: %w", err)
287+
}
288+
opt.ServiceAccountCredentials = string(loadedCreds)
289+
}
290+
291+
if opt.ServiceAccountCredentials != "" {
292+
oAuthClient, err = getServiceAccountClient(ctx, opt, []byte(opt.ServiceAccountCredentials))
293+
if err != nil {
294+
return nil, fmt.Errorf("failed to create oauth client from service account: %w", err)
295+
}
296+
oAuthClients = append(oAuthClients, oAuthClient)
297+
}
298+
299+
// Only create a default oauth client if we don't have one already
300+
if oAuthClient == nil {
301+
oAuthClient, _, err = oauthutil.NewClientWithBaseClient(ctx, name, m, driveConfig, getClient(ctx, opt))
302+
if err != nil {
303+
return nil, fmt.Errorf("failed to create oauth client: %w", err)
304+
}
305+
oAuthClients = append(oAuthClients, oAuthClient)
306+
}
307+
308+
return oAuthClients, nil
309+
}
310+
311+
func (f *Fs) getUploadClient(_ context.Context) *http.Client {
312+
uploadClientMutex := f.uploadClientMutex
313+
uploadClientMutex.Lock()
314+
defer uploadClientMutex.Unlock()
315+
316+
f.uploadClientIndex++
317+
if f.uploadClientIndex > len(f.uploadClients)-1 {
318+
f.uploadClientIndex = 0
319+
}
320+
321+
return f.uploadClients[f.uploadClientIndex]
322+
}

0 commit comments

Comments
 (0)