Skip to content

Commit 24d6af9

Browse files
efiopefiop
committed
fixup
1 parent 571fb3d commit 24d6af9

File tree

1 file changed

+125
-97
lines changed

1 file changed

+125
-97
lines changed

‎pydrive2/fs/spec.py

Lines changed: 125 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import posixpath
77
import threading
88
from collections import defaultdict
9+
from contextlib import contextmanager
910

1011
from fsspec.spec import AbstractFileSystem
1112
from funcy import cached_property, retry, wrap_prop, wrap_with
@@ -20,6 +21,14 @@
2021

2122
FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
2223

24+
COMMON_SETTINGS = {
25+
"get_refresh_token": True,
26+
"oauth_scope": [
27+
"https://www.googleapis.com/auth/drive",
28+
"https://www.googleapis.com/auth/drive.appdata",
29+
],
30+
}
31+
2332

2433
class GDriveAuthError(Exception):
2534
pass
@@ -55,14 +64,91 @@ def should_retry(exc):
5564
)(func)
5665

5766

58-
def _cache_file(client_id, profile=None):
59-
cache_dir = appdirs.user_cache_dir("pydrive2fs", appauthor=False)
60-
client_cache_dir = os.path.join(cache_dir, client_id)
61-
profile = profile or "default"
67+
@contextmanager
68+
def _wrap_errors():
69+
try:
70+
yield
71+
except Exception as exc:
72+
# Handle AuthenticationError, RefreshError and other auth failures
73+
# It's hard to come up with a narrow exception, since PyDrive throws
74+
# a lot of different errors - broken credentials file, refresh token
75+
# expired, flow failed, etc.
76+
raise GDriveAuthError("Failed to authenticate GDrive") from exc
77+
78+
79+
def _client_auth(
80+
client_id=None,
81+
client_secret=None,
82+
client_json=None,
83+
client_json_file_path=None,
84+
profile=None,
85+
):
86+
if client_json:
87+
save_settings = {
88+
"save_credentials_backend": "dictionary",
89+
"save_credentials_dict": {"creds": client_json},
90+
"save_credentials_key": "creds",
91+
}
92+
else:
93+
creds_file = client_json_file_path
94+
if not creds_file:
95+
profile = profile or "default"
96+
cache_dir = os.path.join(
97+
appdirs.user_cache_dir("pydrive2fs", appauthor=False),
98+
client_id,
99+
)
100+
os.makedirs(cache_dir, exist_ok=True)
101+
creds_file = os.path.join(cache_dir, f"{profile}.json")
102+
103+
save_settings = {
104+
"save_credentials_backend": "file",
105+
"save_credentials_file": creds_file,
106+
}
107+
108+
settings = {
109+
**COMMON_SETTINGS,
110+
"save_credentials": True,
111+
**save_settings,
112+
"client_config_backend": "settings",
113+
"client_config": {
114+
"client_id": client_id,
115+
"client_secret": client_secret,
116+
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
117+
"token_uri": "https://oauth2.googleapis.com/token",
118+
"revoke_uri": "https://oauth2.googleapis.com/revoke",
119+
"redirect_uri": "",
120+
},
121+
}
122+
123+
auth = GoogleAuth(settings=settings)
62124

63-
os.makedirs(client_cache_dir, exist_ok=True)
125+
with _wrap_errors():
126+
auth.LocalWebserverAuth()
64127

65-
return os.path.join(client_cache_dir, f"{profile}.json")
128+
return auth
129+
130+
131+
def _service_auth(
132+
client_user_email=None,
133+
client_json=None,
134+
client_json_file_path=None,
135+
):
136+
settings = {
137+
**COMMON_SETTINGS,
138+
"client_config_backend": "service",
139+
"service_config": {
140+
"client_user_email": client_user_email,
141+
"client_json": client_json,
142+
"client_json_file_path": client_json_file_path,
143+
},
144+
}
145+
146+
auth = GoogleAuth(settings=settings)
147+
148+
with _wrap_errors():
149+
auth.ServiceAuth()
150+
151+
return auth
66152

67153

68154
class GDriveFileSystem(AbstractFileSystem):
@@ -71,6 +157,13 @@ def __init__(
71157
path,
72158
google_auth=None,
73159
trash_only=True,
160+
client_id=None,
161+
client_secret=None,
162+
client_user_email=None,
163+
client_json=None,
164+
client_json_file_path=None,
165+
use_service_account=False,
166+
profile=None,
74167
**kwargs,
75168
):
76169
"""Access to gdrive as a file-system
@@ -104,100 +197,35 @@ def __init__(
104197
self.root, self.base = self.split_path(self.path)
105198

106199
if not google_auth:
107-
google_auth = self._auth(**kwargs)
108-
109-
self.client = GoogleDrive(google_auth)
110-
self._trash_only = trash_only
111-
112-
def _auth(
113-
self,
114-
client_id=None,
115-
client_secret=None,
116-
client_user_email=None,
117-
client_json=None,
118-
client_json_file_path=None,
119-
use_service_account=False,
120-
profile=None,
121-
**_kwargs,
122-
):
123-
if (
124-
not client_json
125-
and not client_json_file_path
126-
and not (client_id and client_secret)
127-
):
128-
raise ValueError(
129-
"Specify credentials using one of these methods: "
130-
"client_id/client_secret or "
131-
"client_json or "
132-
"client_json_file_path"
133-
)
134-
135-
common = {
136-
"get_refresh_token": True,
137-
"oauth_scope": [
138-
"https://www.googleapis.com/auth/drive",
139-
"https://www.googleapis.com/auth/drive.appdata",
140-
],
141-
}
142-
143-
if use_service_account:
144-
settings = {
145-
**common,
146-
"client_config_backend": "service",
147-
"service_config": {
148-
"client_user_email": client_user_email,
149-
"client_json": client_json,
150-
"client_json_file_path": client_json_file_path,
151-
},
152-
}
153-
else:
154-
if client_json:
155-
save_settings = {
156-
"save_credentials_backend": "dictionary",
157-
"save_credentials_dict": {"creds": client_json},
158-
"save_credentials_key": "creds",
159-
}
160-
else:
161-
creds_file = client_json_file_path
162-
if not creds_file:
163-
creds_file = _cache_file(client_id, profile)
164-
165-
save_settings = {
166-
"save_credentials_backend": "file",
167-
"save_credentials_file": creds_file,
168-
}
169-
170-
settings = {
171-
**common,
172-
"save_credentials": True,
173-
**save_settings,
174-
"client_config_backend": "settings",
175-
"client_config": {
176-
"client_id": client_id,
177-
"client_secret": client_secret,
178-
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
179-
"token_uri": "https://oauth2.googleapis.com/token",
180-
"revoke_uri": "https://oauth2.googleapis.com/revoke",
181-
"redirect_uri": "",
182-
},
183-
}
184-
185-
google_auth = GoogleAuth(settings=settings)
200+
if (
201+
not client_json
202+
and not client_json_file_path
203+
and not (client_id and client_secret)
204+
):
205+
raise ValueError(
206+
"Specify credentials using one of these methods: "
207+
"client_id/client_secret or "
208+
"client_json or "
209+
"client_json_file_path"
210+
)
186211

187-
try:
188-
logger.debug("GDrive auth with config '%s'.", settings)
189212
if use_service_account:
190-
google_auth.ServiceAuth()
213+
google_auth = _service_auth(
214+
client_json=client_json,
215+
client_json_file_path=client_json_file_path,
216+
client_user_email=client_user_email,
217+
)
191218
else:
192-
google_auth.LocalWebserverAuth()
193-
except Exception as exc:
194-
# Handle AuthenticationError, RefreshError and other auth failures
195-
# It's hard to come up with a narrow exception, since PyDrive throws
196-
# a lot of different errors - broken credentials file, refresh token
197-
# expired, flow failed, etc.
198-
raise GDriveAuthError("Failed to authenticate GDrive") from exc
199-
200-
return google_auth
219+
google_auth = _client_auth(
220+
client_id,
221+
client_secret,
222+
client_json=client_json,
223+
client_json_file_path=client_json_file_path,
224+
profile=profile,
225+
)
226+
227+
self.client = GoogleDrive(google_auth)
228+
self._trash_only = trash_only
201229

202230
def split_path(self, path):
203231
parts = path.replace("//", "/").rstrip("/").split("/", 1)

0 commit comments

Comments
 (0)