outs: support gs/s3 as an output without cache

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

dvc/output/gs.py

@@ -1,5 +1,69 @@
 from dvc.dependency.gs import DependencyGS
 
+
 class OutputGS(DependencyGS):
-    def __init__(self, stage, path, etag=None):
-        raise NotImplemented
+    PARAM_CACHE = 'cache'
+
+    def __init__(self, stage, path, etag=None, use_cache=True):
+        super(OutputGS, self).__init__(stage, path, etag=etag)
+        self.use_cache = use_cache
+
+    @property
+    def cache(self):
+        if not self.use_cache:
+            return None
+
+        raise NotImplementedError
+
+        return self.project.cache.get(self.etag)
+
+    def dumpd(self):
+        ret = super(OutputGS, self).dumpd()
+        ret[self.PARAM_CACHE] = self.use_cache
+        return ret
+
+    @classmethod
+    def loadd(cls, stage, d):
+        ret = super(OutputGS, cls).loadd(stage, d)
+        ret.use_cache = d.get(cls.PARAM_CACHE, True)
+        return ret
+
+    @classmethod
+    def loads(cls, stage, s, use_cache=True):
+        ret = super(OutputGS, cls).loads(stage, s)
+        ret.use_cache = use_cache
+        return ret
+
+    @classmethod
+    def loads_from(cls, stage, s_list, use_cache=True):
+        return [cls.loads(stage, x, use_cache=use_cache) for x in s_list]
+
+    def changed(self):
+        if super(OutputGS, self).changed():
+            return True
+
+        if self.use_cache and self.project.cache.changed(self.etag):
+            return True
+
+        return False
+
+    def checkout(self):
+        if not self.use_cache:
+            return
+
+        raise NotImplementedError
+
+    def save(self):
+        super(OutputGS, self).save()
+
+        if not self.use_cache:
+            return
+
+        raise NotImplementedError
+
+    def remove(self):
+        blob = self.client.bucket(self.bucket).get_blob(self.key)
+        if not blob:
+            return
+
+        blob.delete()

dvc/output/s3.py

@@ -2,5 +2,64 @@ from dvc.dependency.s3 import DependencyS3
 
 
 class OutputS3(DependencyS3):
-    def __init__(self, stage, path, etag=None):
-        raise NotImplemented
+    PARAM_CACHE = 'cache'
+
+    def __init__(self, stage, path, etag=None, use_cache=True):
+        super(OutputS3, self).__init__(stage, path, etag=etag)
+        self.use_cache = use_cache
+
+    @property
+    def cache(self):
+        if not self.use_cache:
+            return None
+
+        raise NotImplementedError
+
+        return self.project.cache.get(self.etag)
+
+    def dumpd(self):
+        ret = super(OutputS3, self).dumpd()
+        ret[self.PARAM_CACHE] = self.use_cache
+        return ret
+
+    @classmethod
+    def loadd(cls, stage, d):
+        ret = super(OutputS3, cls).loadd(stage, d)
+        ret.use_cache = d.get(cls.PARAM_CACHE, True)
+        return ret
+
+    @classmethod
+    def loads(cls, stage, s, use_cache=True):
+        ret = super(OutputS3, cls).loads(stage, s)
+        ret.use_cache = use_cache
+        return ret
+
+    @classmethod
+    def loads_from(cls, stage, s_list, use_cache=True):
+        return [cls.loads(stage, x, use_cache=use_cache) for x in s_list]
+
+    def changed(self):
+        if super(OutputS3, self).changed():
+            return True
+
+        if self.use_cache and self.project.cache.changed(self.etag):
+            return True
+
+        return False
+
+    def checkout(self):
+        if not self.use_cache:
+            return
+
+        raise NotImplementedError
+
+    def save(self):
+        super(OutputS3, self).save()
+
+        if not self.use_cache:
+            return
+
+        raise NotImplementedError
+
+    def remove(self):
+        self.s3.Object(self.bucket, self.key).delete()

tests/test_repro.py

@@ -186,56 +186,74 @@ class TestCmdRepro(TestRepro):
         self.assertNotEqual(ret, 0)
 
 
-class TestReproDependencyS3(TestDvc):
+class TestReproExternalBase(TestDvc):
+    def should_test(self):
+        return False
+
     def test(self):
-        if not _should_test_aws():
+        if not self.should_test():
             return
 
-        key = str(uuid.uuid4()) + '/' + self.FOO
-        path = 's3://' + TEST_AWS_REPO_BUCKET + '/' + key
+        foo_key = str(uuid.uuid4()) + '/' + self.FOO
+        bar_key = str(uuid.uuid4()) + '/' + self.BAR
 
-        s3 = boto3.resource('s3')
+        foo_path = self.schema + self.bucket + '/' + foo_key
+        bar_path = self.schema + self.bucket + '/' + bar_key
 
-        with open(self.FOO, 'rb') as fd:
-            s3.Bucket(TEST_AWS_REPO_BUCKET).put_object(Key=key, Body=fd.read())
+        self.write(self.bucket, foo_key, 'foo')
 
-        stage = self.dvc.run(outs=[self.FOO],
-                             deps=[path],
-                             cmd='aws s3 cp {} {}'.format(path, self.FOO))
+        stage = self.dvc.run(outs_no_cache=[bar_path],
+                             deps=[foo_path],
+                             cmd='{} {} {}'.format(self.cmd, foo_path, bar_path))
 
-        with open(self.BAR, 'rb') as fd:
-            s3.Bucket(TEST_AWS_REPO_BUCKET).put_object(Key=key, Body=fd.read())
+        self.write(self.bucket, foo_key, 'bar')
 
         stages = self.dvc.reproduce(stage.path)
 
         self.assertEqual(len(stages), 1)
-        self.assertTrue(filecmp.cmp(self.FOO, self.BAR))
 
 
-class TestReproDependencyGS(TestDvc):
-    def test(self):
-        if not _should_test_gcp():
-            return
+class TestReproExternalS3(TestReproExternalBase):
+    def should_test(self):
+        return _should_test_aws()
 
+    @property
+    def schema(self):
+        return 's3://'
+
+    @property
+    def bucket(self):
+        return TEST_AWS_REPO_BUCKET
+
+    @property
+    def cmd(self):
+        return 'aws s3 cp'
+
+    def write(self, bucket, key, body):
+        s3 = boto3.resource('s3')
+        s3.Bucket(bucket).put_object(Key=key, Body=body)
+
+
+class TestReproExternalGS(TestReproExternalBase):
+    def should_test(self):
         #FIXME enable
         if os.getenv('CI'):
-            return
+            return False
+        return _should_test_gcp()
 
-        key = str(uuid.uuid4()) + '/' + self.FOO
-        path = 'gs://' + TEST_GCP_REPO_BUCKET + '/' + key
+    @property
+    def schema(self):
+        return 'gs://'
 
+    @property
+    def bucket(self):
+        return TEST_GCP_REPO_BUCKET
+
+    @property
+    def cmd(self):
+        return 'gsutil cp'
+
+    def write(self, bucket, key, body):
         client = gc.Client()
-        bucket = client.bucket(TEST_GCP_REPO_BUCKET)
-
-        bucket.blob(key).upload_from_filename(self.FOO)
-
-        stage = self.dvc.run(outs=[self.FOO],
-                             deps=[path],
-                             cmd='gsutil cp {} {}'.format(path, self.FOO))
-
-        bucket.blob(key).upload_from_filename(self.BAR)
-
-        stages = self.dvc.reproduce(stage.path)
-
-        self.assertEqual(len(stages), 1)
-        self.assertTrue(filecmp.cmp(self.FOO, self.BAR))
+        bucket = client.bucket(bucket)
+        bucket.blob(key).upload_from_string(body)