dvc: cloud: reuse connections when processing targets

Signed-off-by: Ruslan Kuprieiev <[email protected]>

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):
         """

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

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):

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 """

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