Skip to content

Commit aec6c55

Browse files
committed
fs: simplify auth
Compared to how we were using it in dvc before: 1) doesn't require/accept `tmp_dir` at all, caches creds in `appdirs("pydrive2fs", "iterative")` 2) caching is per client_id, so no risk of using mismatched client_id and cached creds 3) doesn't accept `user_credentials_file` (was used as a manual caching destination), because no one really wants to use it anyway. This eliminates the risk of users using mismatched client_id and cached creds. Other than that, this is mostly old code. Fixes #193
1 parent 0f5eb74 commit aec6c55

File tree

3 files changed

+119
-3
lines changed

3 files changed

+119
-3
lines changed

pydrive2/fs/spec.py

Lines changed: 104 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import appdirs
12
import errno
23
import io
34
import logging
@@ -13,12 +14,17 @@
1314

1415
from pydrive2.drive import GoogleDrive
1516
from pydrive2.fs.utils import IterStream
17+
from pydrive2.auth import GoogleAuth
1618

1719
logger = logging.getLogger(__name__)
1820

1921
FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
2022

2123

24+
class GDriveAuthError(Exception):
25+
pass
26+
27+
2228
def _gdrive_retry(func):
2329
def should_retry(exc):
2430
from pydrive2.files import ApiRequestError
@@ -50,12 +56,108 @@ def should_retry(exc):
5056

5157

5258
class GDriveFileSystem(AbstractFileSystem):
53-
def __init__(self, path, google_auth, trash_only=True, **kwargs):
59+
GDRIVE_CREDENTIALS_DATA = "GDRIVE_CREDENTIALS_DATA"
60+
61+
def __init__(
62+
self,
63+
path,
64+
google_auth=None,
65+
trash_only=True,
66+
client_id=None,
67+
client_secret=None,
68+
client_user_email=None,
69+
client_json_file_path=None,
70+
**kwargs,
71+
):
72+
super().__init__(**kwargs)
5473
self.path = path
5574
self.root, self.base = self.split_path(self.path)
75+
76+
self._cache_dir = appdirs.user_cache_dir("pydrive2fs", "iterative")
77+
os.makedirs(self._cache_dir, exist_ok=True)
78+
79+
if not google_auth:
80+
google_auth = self._auth(
81+
client_id=client_id,
82+
client_secret=client_secret,
83+
client_user_email=client_user_email,
84+
client_json_file_path=client_json_file_path,
85+
)
86+
5687
self.client = GoogleDrive(google_auth)
5788
self._trash_only = trash_only
58-
super().__init__(**kwargs)
89+
90+
def _auth(
91+
self,
92+
client_id=None,
93+
client_secret=None,
94+
client_user_email=None,
95+
client_json_file_path=None,
96+
):
97+
settings = {
98+
"get_refresh_token": True,
99+
"oauth_scope": [
100+
"https://www.googleapis.com/auth/drive",
101+
"https://www.googleapis.com/auth/drive.appdata",
102+
],
103+
}
104+
105+
env_creds = os.getenv(self.GDRIVE_CREDENTIALS_DATA)
106+
if env_creds:
107+
settings["save_credentials"] = True
108+
settings["save_credentials_backend"] = "dictionary"
109+
settings["save_credentials_dict"] = {"creds": env_creds}
110+
settings["save_credentials_key"] = "creds"
111+
elif not client_json_file_path:
112+
assert client_id
113+
save_credentials_file = os.path.join(
114+
self._cache_dir, client_id + ".json"
115+
)
116+
117+
settings["save_credentials"] = True
118+
settings["save_credentials_backend"] = "file"
119+
settings["save_credentials_file"] = save_credentials_file
120+
121+
if client_json_file_path:
122+
config = {}
123+
124+
if client_user_email:
125+
config["client_user_email"] = client_user_email
126+
127+
if env_creds:
128+
config["client_json"] = env_creds
129+
else:
130+
config["client_json_file_path"] = client_json_file_path
131+
132+
settings["client_config_backend"] = "service"
133+
settings["service_config"] = config
134+
else:
135+
settings["client_config_backend"] = "settings"
136+
settings["client_config"] = {
137+
"client_id": client_id,
138+
"client_secret": client_secret,
139+
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
140+
"token_uri": "https://oauth2.googleapis.com/token",
141+
"revoke_uri": "https://oauth2.googleapis.com/revoke",
142+
"redirect_uri": "",
143+
}
144+
145+
google_auth = GoogleAuth(settings=settings)
146+
147+
try:
148+
logger.debug("GDrive auth with config '%s'.", settings)
149+
if client_json_file_path:
150+
google_auth.ServiceAuth()
151+
else:
152+
google_auth.LocalWebserverAuth()
153+
except Exception as exc:
154+
# Handle AuthenticationError, RefreshError and other auth failures
155+
# It's hard to come up with a narrow exception, since PyDrive throws
156+
# a lot of different errors - broken credentials file, refresh token
157+
# expired, flow failed, etc.
158+
raise GDriveAuthError("Failed to authenticate GDrive") from exc
159+
160+
return google_auth
59161

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

pydrive2/test/test_fs.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,15 @@ def fs(tmpdir, base_remote_dir):
3636
return fs
3737

3838

39+
def test_fs_service_from_settings(base_remote_dir):
40+
creds = "credentials/fs.dat"
41+
setup_credentials(creds)
42+
GDriveFileSystem(
43+
base_remote_dir,
44+
service_account_json_file_path=creds,
45+
)
46+
47+
3948
def test_info(fs, tmpdir, remote_dir):
4049
fs.touch(remote_dir + "/info/a.txt")
4150
fs.touch(remote_dir + "/info/b.txt")

setup.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,12 @@
3737
"pyOpenSSL >= 19.1.0",
3838
],
3939
extras_require={
40-
"fsspec": ["fsspec >= 2021.07.0", "tqdm >= 4.0.0", "funcy >= 1.14"],
40+
"fsspec": [
41+
"fsspec >= 2021.07.0",
42+
"tqdm >= 4.0.0",
43+
"funcy >= 1.14",
44+
"appdirs",
45+
],
4146
"tests": tests_requirements,
4247
},
4348
python_requires=">=3.7",

0 commit comments

Comments
 (0)