Skip to content

Commit 9f6d469

Browse files
maxhoramaxhora
committed
gdrive lazy init
1 parent b883bb9 commit 9f6d469

File tree

1 file changed

+80
-64
lines changed

1 file changed

+80
-64
lines changed

‎dvc/remote/gdrive/__init__.py

Lines changed: 80 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33
import os
44
import posixpath
55
import logging
6+
import threading
67

7-
from funcy import cached_property, retry, compose, decorator
8+
from funcy import cached_property, retry, compose, decorator, wrap_with
89
from funcy.py3 import cat
910

1011
from dvc.remote.gdrive.utils import TrackFileReadProgress, FOLDER_MIME_TYPE
@@ -84,9 +85,6 @@ def init_drive(self):
8485
)
8586
)
8687

87-
self.root_id = self.get_remote_id(self.path_info, create=True)
88-
self.cached_dirs, self.cached_ids = self.cache_root_dirs()
89-
9088
def gdrive_upload_file(
9189
self, args, no_progress_bar=True, from_file="", progress_name=""
9290
):
@@ -132,48 +130,66 @@ def cache_root_dirs(self):
132130
cached_ids[dir1["id"]] = dir1["title"]
133131
return cached_dirs, cached_ids
134132

135-
@cached_property
133+
@property
134+
def cached_dirs(self):
135+
if not hasattr(self, '_cached_dirs'):
136+
self.drive
137+
return self._cached_dirs
138+
139+
@property
140+
def cached_ids(self):
141+
if not hasattr(self, '_cached_ids'):
142+
self.drive
143+
return self._cached_ids
144+
145+
@property
146+
@wrap_with(threading.RLock())
136147
def drive(self):
137-
from pydrive.auth import GoogleAuth
138-
from pydrive.drive import GoogleDrive
139-
140-
if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
141-
with open(
142-
self.gdrive_user_credentials_path, "w"
143-
) as credentials_file:
144-
credentials_file.write(
145-
os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA)
146-
)
147-
148-
GoogleAuth.DEFAULT_SETTINGS["client_config_backend"] = "settings"
149-
GoogleAuth.DEFAULT_SETTINGS["client_config"] = {
150-
"client_id": self.gdrive_client_id,
151-
"client_secret": self.gdrive_client_secret,
152-
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
153-
"token_uri": "https://oauth2.googleapis.com/token",
154-
"revoke_uri": "https://oauth2.googleapis.com/revoke",
155-
"redirect_uri": "",
156-
}
157-
GoogleAuth.DEFAULT_SETTINGS["save_credentials"] = True
158-
GoogleAuth.DEFAULT_SETTINGS["save_credentials_backend"] = "file"
159-
GoogleAuth.DEFAULT_SETTINGS[
160-
"save_credentials_file"
161-
] = self.gdrive_user_credentials_path
162-
GoogleAuth.DEFAULT_SETTINGS["get_refresh_token"] = True
163-
GoogleAuth.DEFAULT_SETTINGS["oauth_scope"] = [
164-
"https://www.googleapis.com/auth/drive",
165-
"https://www.googleapis.com/auth/drive.appdata",
166-
]
167-
168-
# Pass non existent settings path to force DEFAULT_SETTINGS loading
169-
gauth = GoogleAuth(settings_file="")
170-
gauth.CommandLineAuth()
171-
172-
if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
173-
os.remove(self.gdrive_user_credentials_path)
174-
175-
gdrive = GoogleDrive(gauth)
176-
return gdrive
148+
if not hasattr(self, '_gdrive'):
149+
from pydrive.auth import GoogleAuth
150+
from pydrive.drive import GoogleDrive
151+
152+
if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
153+
with open(
154+
self.gdrive_user_credentials_path, "w"
155+
) as credentials_file:
156+
credentials_file.write(
157+
os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA)
158+
)
159+
160+
GoogleAuth.DEFAULT_SETTINGS["client_config_backend"] = "settings"
161+
GoogleAuth.DEFAULT_SETTINGS["client_config"] = {
162+
"client_id": self.gdrive_client_id,
163+
"client_secret": self.gdrive_client_secret,
164+
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
165+
"token_uri": "https://oauth2.googleapis.com/token",
166+
"revoke_uri": "https://oauth2.googleapis.com/revoke",
167+
"redirect_uri": "",
168+
}
169+
GoogleAuth.DEFAULT_SETTINGS["save_credentials"] = True
170+
GoogleAuth.DEFAULT_SETTINGS["save_credentials_backend"] = "file"
171+
GoogleAuth.DEFAULT_SETTINGS[
172+
"save_credentials_file"
173+
] = self.gdrive_user_credentials_path
174+
GoogleAuth.DEFAULT_SETTINGS["get_refresh_token"] = True
175+
GoogleAuth.DEFAULT_SETTINGS["oauth_scope"] = [
176+
"https://www.googleapis.com/auth/drive",
177+
"https://www.googleapis.com/auth/drive.appdata",
178+
]
179+
180+
# Pass non existent settings path to force DEFAULT_SETTINGS loading
181+
gauth = GoogleAuth(settings_file="")
182+
gauth.CommandLineAuth()
183+
184+
if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA):
185+
os.remove(self.gdrive_user_credentials_path)
186+
187+
self._gdrive = GoogleDrive(gauth)
188+
189+
self.root_id = self.get_remote_id(self.path_info, create=True)
190+
self._cached_dirs, self._cached_ids = self.cache_root_dirs()
191+
192+
return self._gdrive
177193

178194
@gdrive_retry
179195
def create_remote_dir(self, parent_id, title):
@@ -211,39 +227,39 @@ def resolve_remote_item_from_path(self, parents_ids, path_parts, create):
211227
parents_ids = [item["id"]]
212228
return item
213229

214-
def subtract_root_path(self, parts):
230+
def subtract_root_path(self, path_parts):
215231
if not hasattr(self, "root_id"):
216-
return parts, [self.path_info.bucket]
232+
return path_parts, [self.path_info.bucket]
217233

218234
for part in self.path_info.path.split("/"):
219-
if parts and parts[0] == part:
220-
parts.pop(0)
235+
if path_parts and path_parts[0] == part:
236+
path_parts.pop(0)
221237
else:
222238
break
223-
return parts, [self.root_id]
239+
return path_parts, [self.root_id]
224240

225241
def get_remote_id_from_cache(self, path_info):
226-
files_ids = []
227-
parts, parents_ids = self.subtract_root_path(path_info.path.split("/"))
242+
remote_ids = []
243+
path_parts, parents_ids = self.subtract_root_path(path_info.path.split("/"))
228244
if (
229-
hasattr(self, "cached_dirs")
245+
hasattr(self, "_cached_dirs")
230246
and path_info != self.path_info
231-
and parts
232-
and (parts[0] in self.cached_dirs)
247+
and path_parts
248+
and (path_parts[0] in self.cached_dirs)
233249
):
234-
parents_ids = self.cached_dirs[parts[0]]
235-
files_ids = self.cached_dirs[parts[0]]
236-
parts.pop(0)
250+
parents_ids = self.cached_dirs[path_parts[0]]
251+
remote_ids = self.cached_dirs[path_parts[0]]
252+
path_parts.pop(0)
237253

238-
return files_ids, parents_ids, parts
254+
return remote_ids, parents_ids, path_parts
239255

240256
def get_remote_id(self, path_info, create=False):
241-
files_ids, parents_ids, parts = self.get_remote_id_from_cache(path_info)
257+
remote_ids, parents_ids, path_parts = self.get_remote_id_from_cache(path_info)
242258

243-
if not parts and files_ids:
244-
return files_ids[0]
259+
if not path_parts and remote_ids:
260+
return remote_ids[0]
245261

246-
file1 = self.resolve_remote_item_from_path(parents_ids, parts, create)
262+
file1 = self.resolve_remote_item_from_path(parents_ids, path_parts, create)
247263
return file1["id"] if file1 else ""
248264

249265
def exists(self, path_info):
@@ -294,7 +310,7 @@ def list_remote_item(self, drive_file):
294310
yield drive_file["title"]
295311

296312
def all(self):
297-
if not hasattr(self, "cached_ids") or not self.cached_ids:
313+
if not self.cached_ids:
298314
return
299315

300316
query = " or ".join(

0 commit comments

Comments
 (0)