gdrive: support service auth with json (#5485)

* gdrive: support service auth with json

* retry logic

.github/workflows/tests.yaml

@@ -27,7 +27,11 @@ jobs:
       with:
         python-version: 3.8
     - name: Install requirements
-      run: pip install ".[all,tests]" Pygments collective.checkdocs pre-commit git+https://github.com/intake/gdrivefs.git
+      run: pip install ".[all,tests]" Pygments collective.checkdocs pre-commit
+    - name: install temporary dependencies
+      run: |
+        pip install git+https://github.com/isidentical/gdrivefs@service-account
+        pip install git+https://github.com/isidentical/pydata-google-auth@service-account
     - name: Check README
       run: python setup.py checkdocs
     - uses: pre-commit/[email protected]
@@ -69,6 +73,10 @@ jobs:
       run: |
         pip install wheel
         pip install -e ".[all,tests]"
+    - name: install temporary dependencies
+      run: |
+        pip install git+https://github.com/isidentical/gdrivefs@service-account
+        pip install git+https://github.com/isidentical/pydata-google-auth@service-account
     - name: setup git
       run: |
         git config --global user.email "[email protected]"

dvc/config.py

@@ -236,7 +236,7 @@ class Config(dict):
                 str: {
                     "url": func,
                     "gdrive_user_credentials_file": func,
-                    "gdrive_service_account_p12_file_path": func,
+                    "gdrive_service_account_json_file_path": func,
                     "credentialpath": func,
                     "keyfile": func,
                     "cert_path": func,

dvc/config_schema.py

@@ -191,9 +191,8 @@ SCHEMA = {
                     "gdrive_client_id": str,
                     "gdrive_client_secret": str,
                     "gdrive_user_credentials_file": str,
-                    "gdrive_service_account_email": str,
                     "gdrive_service_account_user_email": str,
-                    "gdrive_service_account_p12_file_path": str,
+                    "gdrive_service_account_json_file_path": str,
                     Optional("gdrive_trash_only", default=False): Bool,
                     **REMOTE_COMMON,
                 },

dvc/fs/gdrive.py

@@ -117,14 +117,11 @@ class GDriveFileSystem(BaseFileSystem):
         self._path = self.path_info.path
         self._trash_only = config.get("gdrive_trash_only")
         self._use_service_account = config.get("gdrive_use_service_account")
-        self._service_account_email = config.get(
-            "gdrive_service_account_email"
-        )
         self._service_account_user_email = config.get(
             "gdrive_service_account_user_email"
         )
-        self._service_account_p12_file_path = config.get(
-            "gdrive_service_account_p12_file_path"
+        self._service_account_json_file_path = config.get(
+            "gdrive_service_account_json_file_path"
         )
         self._client_id = config.get("gdrive_client_id")
         self._client_secret = config.get("gdrive_client_secret")
@@ -142,14 +139,14 @@ class GDriveFileSystem(BaseFileSystem):
 
     def _validate_config(self):
         # Validate Service Account configuration
-        if self._use_service_account and (
-            not self._service_account_email
-            or not self._service_account_p12_file_path
+        if (
+            self._use_service_account
+            and not self._service_account_json_file_path
         ):
             raise DvcException(
-                "To use service account, set `gdrive_service_account_email`,\n"
-                "`gdrive_service_account_p12_file_path`, and optionally "
-                "`gdrive_service_account_user_email`\nin DVC config. "
+                "To use service account, set "
+                "`gdrive_service_account_json_file_path`, and optionally"
+                "`gdrive_service_account_user_email` in DVC config\n"
                 "Learn more at {}".format(
                     format_link("https://man.dvc.org/remote/modify")
                 )
@@ -220,7 +217,6 @@ class GDriveFileSystem(BaseFileSystem):
         auth_settings = {
             "client_config_backend": "settings",
             "save_credentials": True,
-            "save_credentials_backend": "file",
             "save_credentials_file": self._gdrive_user_credentials_path,
             "get_refresh_token": True,
             "oauth_scope": [
@@ -230,12 +226,13 @@ class GDriveFileSystem(BaseFileSystem):
         }
 
         if self._use_service_account:
+            auth_settings["save_credentials_backend"] = "json"
             auth_settings["service_config"] = {
-                "client_service_email": self._service_account_email,
                 "client_user_email": self._service_account_user_email,
-                "client_pkcs12_file_path": self._service_account_p12_file_path,
+                "client_json_file_path": self._service_account_json_file_path,
             }
         else:
+            auth_settings["save_credentials_backend"] = "file"
             auth_settings["client_config"] = {
                 "client_id": self._client_id or self.DEFAULT_GDRIVE_CLIENT_ID,
                 "client_secret": self._client_secret

setup.py

@@ -94,7 +94,7 @@ install_requires = [
 # Extra dependencies for remote integrations
 
 gs = ["google-cloud-storage==1.19.0"]
-gdrive = ["pydrive2>=1.6.3", "six >= 1.13.0"]
+gdrive = ["pydrive2>=1.7.3", "six >= 1.13.0"]
 s3 = ["boto3>=1.9.201"]
 azure = ["azure-storage-blob>=12.0", "knack"]
 # https://github.com/Legrandin/pycryptodome/issues/465

tests/func/test_config.py

@@ -218,7 +218,7 @@ def test_config_loads_without_error_for_non_dvc_repo(tmp_dir):
         ("keyfile", "ssh://[email protected]:1234/path/to/dir"),
         ("cert_path", "webdavs://example.com/files/USERNAME/"),
         ("key_path", "webdavs://example.com/files/USERNAME/"),
-        ("gdrive_service_account_p12_file_path", "gdrive://root/test"),
+        ("gdrive_service_account_json_file_path", "gdrive://root/test"),
         ("gdrive_user_credentials_file", "gdrive://root/test"),
     ],
 )

tests/remotes/gdrive.py

@@ -1,10 +1,11 @@
 # pylint:disable=abstract-method
+import json
 import os
 import uuid
-from functools import partial, partialmethod
+from functools import partialmethod
 
 import pytest
-from funcy import cached_property
+from funcy import cached_property, retry
 
 from dvc.fs.gdrive import GDriveFileSystem
 from dvc.path_info import CloudURLInfo
@@ -14,10 +15,36 @@ from .base import Base
 
 TEST_GDRIVE_REPO_BUCKET = "root"
 
-CREDENTIALS_DIR = os.path.join(os.path.expanduser("~"), ".config", "pydata")
-CREDENTIALS_FILE = os.path.join(
-    CREDENTIALS_DIR, "pydata_google_credentials.json"
-)
+
+def _gdrive_retry(func):
+    def should_retry(exc):
+        from googleapiclient.errors import HttpError
+
+        if not isinstance(exc, HttpError):
+            return False
+
+        if 500 <= exc.resp.status < 600:
+            return True
+
+        if exc.resp.status == 403:
+            try:
+                reason = json.loads(exc.content)["error"]["errors"][0][
+                    "reason"
+                ]
+            except (ValueError, LookupError):
+                return False
+
+            return reason in [
+                "userRateLimitExceeded",
+                "rateLimitExceeded",
+            ]
+
+    # 16 tries, start at 0.5s, multiply by golden ratio, cap at 20s
+    return retry(
+        16,
+        timeout=lambda a: min(0.5 * 1.618 ** a, 20),
+        filter_errors=should_retry,
+    )(func)
 
 
 class GDrive(Base, CloudURLInfo):
@@ -27,11 +54,23 @@ class GDrive(Base, CloudURLInfo):
 
     @cached_property
     def config(self):
+        tmp_path = tmp_fname()
+        with open(tmp_path, "w") as stream:
+            raw_credentials = os.getenv(
+                GDriveFileSystem.GDRIVE_CREDENTIALS_DATA
+            )
+            try:
+                credentials = json.loads(raw_credentials)
+            except ValueError:
+                credentials = {}
+
+            use_service_account = credentials.get("type") == "service_account"
+            stream.write(raw_credentials)
+
         return {
             "url": self.url,
-            "gdrive_service_account_email": "test",
-            "gdrive_service_account_p12_file_path": "test.p12",
-            "gdrive_use_service_account": True,
+            "gdrive_service_account_json_file_path": tmp_path,
+            "gdrive_use_service_account": use_service_account,
         }
 
     @staticmethod
@@ -45,45 +84,38 @@ class GDrive(Base, CloudURLInfo):
 
     @cached_property
     def client(self):
-        import pydata_google_auth
-
         try:
             from gdrivefs import GoogleDriveFileSystem
         except ImportError:
             pytest.skip("gdrivefs is not installed")
 
-        tmp_path = tmp_fname()
-        with open(tmp_path, "w") as stream:
-            stream.write(os.getenv(GDriveFileSystem.GDRIVE_CREDENTIALS_DATA))
-
-        GoogleDriveFileSystem._connect_cache = partial(
-            pydata_google_auth.load_user_credentials, tmp_path
+        return GoogleDriveFileSystem(
+            token="cache",
+            tokens_file=self.config["gdrive_service_account_json_file_path"],
+            service_account=self.config["gdrive_use_service_account"],
         )
-        return GoogleDriveFileSystem(token="cache")
 
+    @_gdrive_retry
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         if not self.client.exists(self.path):
             self.client.mkdir(self.path)
 
+    @_gdrive_retry
     def write_bytes(self, contents):
         with self.client.open(self.path, mode="wb") as stream:
             stream.write(contents)
 
+    @_gdrive_retry
     def _read(self, mode):
         with self.client.open(self.path, mode=mode) as stream:
             return stream.read()
 
-    def cleanup(self):
-        self.client.delete(self.path, recursive=True)
-
     read_text = partialmethod(_read, mode="r")
     read_bytes = partialmethod(_read, mode="rb")
 
 
 @pytest.fixture
 def gdrive(test_config, make_tmp_dir):
-    pytest.skip("temporarily disabled")
-
     test_config.requires("gdrive")
     if not GDrive.should_test():
         pytest.skip("no gdrive")