|
| 1 | +import appdirs |
1 | 2 | import errno
|
2 | 3 | import io
|
3 | 4 | import logging
|
|
13 | 14 |
|
14 | 15 | from pydrive2.drive import GoogleDrive
|
15 | 16 | from pydrive2.fs.utils import IterStream
|
| 17 | +from pydrive2.auth import GoogleAuth |
16 | 18 |
|
17 | 19 | logger = logging.getLogger(__name__)
|
18 | 20 |
|
19 | 21 | FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
|
20 | 22 |
|
21 | 23 |
|
| 24 | +class GDriveAuthError(Exception): |
| 25 | + pass |
| 26 | + |
| 27 | + |
22 | 28 | def _gdrive_retry(func):
|
23 | 29 | def should_retry(exc):
|
24 | 30 | from pydrive2.files import ApiRequestError
|
@@ -50,12 +56,108 @@ def should_retry(exc):
|
50 | 56 |
|
51 | 57 |
|
52 | 58 | 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) |
54 | 73 | self.path = path
|
55 | 74 | 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 | + |
56 | 87 | self.client = GoogleDrive(google_auth)
|
57 | 88 | 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 |
59 | 161 |
|
60 | 162 | def split_path(self, path):
|
61 | 163 | parts = path.replace("//", "/").rstrip("/").split("/", 1)
|
|
0 commit comments