Skip to content

Commit

Permalink
gdrive: add support for teamDriveId
Browse files Browse the repository at this point in the history
  • Loading branch information
maxhora committed Jan 11, 2020
1 parent b7f9e73 commit 57111a7
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 7 deletions.
19 changes: 16 additions & 3 deletions dvc/remote/gdrive.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import posixpath
import logging
import threading
import re

from funcy import retry, compose, decorator, wrap_with
from funcy.py3 import cat
Expand Down Expand Up @@ -68,6 +69,16 @@ def __init__(self, repo, config):
super().__init__(repo, config)
self.no_traverse = False
self.path_info = self.path_cls(config[Config.SECTION_REMOTE_URL])

bucket = re.search(
"{}://(.*)".format(self.scheme),
config[Config.SECTION_REMOTE_URL],
re.IGNORECASE,
)
self.bucket = (
bucket.group(1).split("/")[0] if bucket else self.path_info.bucket
)

self.config = config
self.init_drive()

Expand Down Expand Up @@ -144,7 +155,7 @@ def cache_root_dirs(self):
cached_dirs = {}
cached_ids = {}
for dir1 in self.gdrive_list_item(
"'{}' in parents and trashed=false".format(self.root_id)
"'{}' in parents and trashed=false".format(self.remote_root_id)
):
remote_path = posixpath.join(self.path_info.path, dir1["title"])
cached_dirs.setdefault(remote_path, []).append(dir1["id"])
Expand Down Expand Up @@ -227,7 +238,9 @@ def drive(self):

self._gdrive = GoogleDrive(gauth)

self.root_id = self.get_remote_id(self.path_info, create=True)
self.remote_root_id = self.get_remote_id(
self.path_info, create=True
)
self._cached_dirs, self._cached_ids = self.cache_root_dirs()

return self._gdrive
Expand Down Expand Up @@ -261,7 +274,7 @@ def get_remote_item(self, name, parents_ids):
return next(iter(item_list), None)

def resolve_remote_item_from_path(self, path_parts, create):
parents_ids = ["root"]
parents_ids = [self.bucket]
current_path = ""
for path_part in path_parts:
current_path = posixpath.join(current_path, path_part)
Expand Down
5 changes: 4 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,10 @@ def run(self):
# google-api-python-client is internal dependency of pydrive. After merge of
# https://github.com/gsuitedevs/PyDrive/pull/180 into pydrive's master,
# usage of google-api-python-client can be removed from DVC.
gdrive = ["pydrive==1.3.1", "google-api-python-client>=1.2"]
gdrive = [
"pydrive @ git+https://github.com/gsuitedevs/PyDrive@master#egg=pydrive",
"google-api-python-client>=1.2",
]
s3 = ["boto3>=1.9.201"]
azure = ["azure-storage-blob==2.1.0"]
oss = ["oss2==2.6.1"]
Expand Down
4 changes: 1 addition & 3 deletions tests/unit/remote/test_gdrive.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
from unittest import TestCase

import pytest
import os

Expand All @@ -18,7 +16,7 @@ class Repo(object):
tmp_dir = ""


class TestRemoteGDrive(TestCase):
class TestRemoteGDrive(object):
CONFIG = {
"url": "gdrive://root/data",
"gdrive_client_id": "client",
Expand Down

0 comments on commit 57111a7

Please sign in to comment.