Skip to content

Commit 3d4a397

Browse files
efiopefiop
committed
fixup
1 parent 571fb3d commit 3d4a397

File tree

1 file changed

+128
-99
lines changed

1 file changed

+128
-99
lines changed

‎pydrive2/fs/spec.py

Lines changed: 128 additions & 99 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,92 @@ 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,
81+
client_secret,
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+
cache_dir = os.path.join(
96+
appdirs.user_cache_dir("pydrive2fs", appauthor=False),
97+
client_id,
98+
)
99+
os.makedirs(cache_dir, exist_ok=True)
100+
101+
profile = profile or "default"
102+
creds_file = os.path.join(cache_dir, f"{profile}.json")
103+
104+
save_settings = {
105+
"save_credentials_backend": "file",
106+
"save_credentials_file": creds_file,
107+
}
108+
109+
settings = {
110+
**COMMON_SETTINGS,
111+
"save_credentials": True,
112+
**save_settings,
113+
"client_config_backend": "settings",
114+
"client_config": {
115+
"client_id": client_id,
116+
"client_secret": client_secret,
117+
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
118+
"token_uri": "https://oauth2.googleapis.com/token",
119+
"revoke_uri": "https://oauth2.googleapis.com/revoke",
120+
"redirect_uri": "",
121+
},
122+
}
123+
124+
auth = GoogleAuth(settings=settings)
125+
126+
with _wrap_errors():
127+
auth.LocalWebserverAuth()
128+
129+
return auth
130+
131+
132+
def _service_auth(
133+
client_user_email=None,
134+
client_json=None,
135+
client_json_file_path=None,
136+
):
137+
settings = {
138+
**COMMON_SETTINGS,
139+
"client_config_backend": "service",
140+
"service_config": {
141+
"client_user_email": client_user_email,
142+
"client_json": client_json,
143+
"client_json_file_path": client_json_file_path,
144+
},
145+
}
146+
147+
auth = GoogleAuth(settings=settings)
62148

63-
os.makedirs(client_cache_dir, exist_ok=True)
149+
with _wrap_errors():
150+
auth.ServiceAuth()
64151

65-
return os.path.join(client_cache_dir, f"{profile}.json")
152+
return auth
66153

67154

68155
class GDriveFileSystem(AbstractFileSystem):
@@ -71,6 +158,13 @@ def __init__(
71158
path,
72159
google_auth=None,
73160
trash_only=True,
161+
client_id=None,
162+
client_secret=None,
163+
client_user_email=None,
164+
client_json=None,
165+
client_json_file_path=None,
166+
use_service_account=False,
167+
profile=None,
74168
**kwargs,
75169
):
76170
"""Access to gdrive as a file-system
@@ -86,7 +180,7 @@ def __init__(
86180
:param client_secret: Client secret of the application.
87181
:type client_secret: str.
88182
:param client_user_email: User email that authority was delegated to
89-
(only for service account).
183+
(only for service account).
90184
:type client_user_email: str.
91185
:param client_json: JSON keyfile loaded into a string.
92186
:type client_json: str.
@@ -95,7 +189,7 @@ def __init__(
95189
:param use_service_account: Use service account.
96190
:type use_service_account: bool.
97191
:param profile: Profile name for caching credentials
98-
(ignored for service account).
192+
(ignored for service account).
99193
:type profile: str.
100194
:raises: GDriveAuthError
101195
"""
@@ -104,100 +198,35 @@ def __init__(
104198
self.root, self.base = self.split_path(self.path)
105199

106200
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)
201+
if (
202+
not client_json
203+
and not client_json_file_path
204+
and not (client_id and client_secret)
205+
):
206+
raise ValueError(
207+
"Specify credentials using one of these methods: "
208+
"client_id/client_secret or "
209+
"client_json or "
210+
"client_json_file_path"
211+
)
186212

187-
try:
188-
logger.debug("GDrive auth with config '%s'.", settings)
189213
if use_service_account:
190-
google_auth.ServiceAuth()
214+
google_auth = _service_auth(
215+
client_json=client_json,
216+
client_json_file_path=client_json_file_path,
217+
client_user_email=client_user_email,
218+
)
191219
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
220+
google_auth = _client_auth(
221+
client_id=client_id,
222+
client_secret=client_secret,
223+
client_json=client_json,
224+
client_json_file_path=client_json_file_path,
225+
profile=profile,
226+
)
227+
228+
self.client = GoogleDrive(google_auth)
229+
self._trash_only = trash_only
201230

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

0 commit comments

Comments
 (0)