diff --git a/dvc/output/gs.py b/dvc/output/gs.py index 0cfa59b9c..e70e9ab79 100644 --- a/dvc/output/gs.py +++ b/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() diff --git a/dvc/output/s3.py b/dvc/output/s3.py index b0971dc66..84cf36a52 100644 --- a/dvc/output/s3.py +++ b/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() diff --git a/tests/test_repro.py b/tests/test_repro.py index 5c13660e0..b5f5f9cd8 100644 --- a/tests/test_repro.py +++ b/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)