diff --git a/dvc/cloud/aws.py b/dvc/cloud/aws.py index af7ccb506..70beddf74 100644 --- a/dvc/cloud/aws.py +++ b/dvc/cloud/aws.py @@ -82,8 +82,7 @@ class DataCloudAWS(DataCloudBase): paths.append(default) return paths - def _get_bucket_aws(self, bucket_name): - """ get a bucket object, aws """ + def connect(self): if all([self._aws_creds.access_key_id, self._aws_creds.secret_access_key, self.aws_region_host]): @@ -92,10 +91,9 @@ class DataCloudAWS(DataCloudBase): host=self.aws_region_host) else: conn = S3Connection() - bucket = conn.lookup(bucket_name) - if bucket is None: - raise DataCloudError('Storage path {} is not setup correctly'.format(bucket_name)) - return bucket + self.bucket = conn.lookup(self.storage_bucket) + if self.bucket is None: + raise DataCloudError('Storage path {} is not setup correctly'.format(self.storage_bucket)) @staticmethod def _cmp_checksum(key, fname): @@ -166,21 +164,18 @@ class DataCloudAWS(DataCloudBase): def _get_key(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_aws(self.storage_bucket) - return bucket.get_key(key_name) + return self.bucket.get_key(key_name) def _get_keys(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_aws(self.storage_bucket) - keys = bucket.list(prefix=key_name) + keys = self.bucket.list(prefix=key_name) if not keys: return None return list(keys) def _new_key(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_aws(self.storage_bucket) - return bucket.new_key(key_name) + return self.bucket.new_key(key_name) def _write_upload_tracker(self, fname, mp_id): """ diff --git a/dvc/cloud/base.py b/dvc/cloud/base.py index cc5ea654e..2bbb2beed 100644 --- a/dvc/cloud/base.py +++ b/dvc/cloud/base.py @@ -207,3 +207,6 @@ class DataCloudBase(object): res.append(self._status(cache_key, cache)) return max(res) + + def connect(self): + pass diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index f97c701ac..5ce068833 100644 --- a/dvc/cloud/data_cloud.py +++ b/dvc/cloud/data_cloud.py @@ -83,6 +83,7 @@ class DataCloud(object): """ Process targets as data items in parallel. """ + self._cloud.connect() return map_progress(func, targets, jobs) def push(self, targets, jobs=1): diff --git a/dvc/cloud/gcp.py b/dvc/cloud/gcp.py index 5a36185cf..c6595f82f 100644 --- a/dvc/cloud/gcp.py +++ b/dvc/cloud/gcp.py @@ -25,14 +25,12 @@ class DataCloudGCP(DataCloudBase): raise ConfigError('can\'t read google cloud project name. ' 'Please set ProjectName in section GC.') - def _get_bucket_gc(self, storage_bucket): - """ get a bucket object, gc """ + def connect(self): client = gc.Client(project=self.gc_project_name) - bucket = client.bucket(storage_bucket) - if not bucket.exists(): + self.bucket = client.bucket(self.storage_bucket) + if not self.bucket.exists(): raise DataCloudError('sync up: google cloud bucket {} ' 'doesn\'t exist'.format(self.storage_bucket)) - return bucket @staticmethod def _cmp_checksum(blob, fname): @@ -79,21 +77,18 @@ class DataCloudGCP(DataCloudBase): def _get_key(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_gc(self.storage_bucket) - return bucket.get_blob(key_name) + return self.bucket.get_blob(key_name) def _get_keys(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_gc(self.storage_bucket) - keys = bucket.list_blobs(prefix=key_name) + keys = self.bucket.list_blobs(prefix=key_name) if not keys: return None return list(keys) def _new_key(self, path): key_name = self.cache_file_key(path) - bucket = self._get_bucket_gc(self.storage_bucket) - return bucket.blob(key_name) + return self.bucket.blob(key_name) def _push_key(self, key, path): """ push, gcp version """ diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 87494f269..5021768e4 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -33,6 +33,8 @@ class TestDataCloudBase(TestDvc): def _test_cloud(self): self._setup_cloud() + self.cloud.connect() + stage = self.dvc.add(self.FOO) cache = stage.outs[0].cache