Repository: iterative/dvc
Pull request #868 (2018-07-09 19:48:25)

Add Azure Blob Storage remote
This change adds a new cloud remote based on [Azure Blob Storage](https://azure.microsoft.com/en-us/services/storage/blobs/) and the [azure-storage-blob SDK](https://github.com/Azure/azure-storage-python).

The implementation closely follows the patterns laid out by the existing cloud remotes (AWS and GCP). Notable differences include:

- No use of a lock in the progress-bar callback and no requirement to look up total transfer size ahead of time: the Azure Storage SDK provides the total transfer bytes and current transfer offset in its progress callback.

- When `exists` is called, check the existence of every blob instead of listing the entire container. Container list operations can be expensive so unless the number of checked files is very large, making separate requests is likely going to be more efficient. *Please let me know if a list-then-check implementation of the exists function would be preferred.*

- For increased flexibility, access credentials for the Azure Storage backend can be provided either connection-string style via the remote URL or via environment variables. The Azure Storage backend can be emulated via [Azurite](https://github.com/Azure/azurite).

This change was created together with @EricSchles.


Patch:
diff --git a/README.rst b/README.rst
index 43a68d6a2..92aeb751b 100644
--- a/README.rst
+++ b/README.rst
@@ -94,7 +94,7 @@ Related technologies
 #. `Workflow Management Systems <https://en.wikipedia.org/wiki/Workflow_management_system>`_. DVC is workflow management system designed specificaly to manage machine learning experiments. DVC was built on top of Git.
 
 DVC is compatible with Git for storing code and the dependency graph (DAG), but not data files cache.
-Data files caches can be transferred separately - now data cache transfer throught AWS S3 and GCP storge are supported.
+Data files caches can be transferred separately - now data cache transfer throught AWS S3, Azure Blob Storage and GCP storge are supported.
 
 How DVC works
 =============
diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py
index e608da938..4a163fb68 100644
--- a/dvc/data_cloud.py
+++ b/dvc/data_cloud.py
@@ -7,6 +7,7 @@ from dvc.config import Config, ConfigError
 
 from dvc.remote.s3 import RemoteS3
 from dvc.remote.gs import RemoteGS
+from dvc.remote.azure import RemoteAzure
 from dvc.remote.ssh import RemoteSSH
 from dvc.remote.hdfs import RemoteHDFS
 from dvc.remote.local import RemoteLOCAL
@@ -19,6 +20,7 @@ class DataCloud(object):
     CLOUD_MAP = {
         'aws'   : RemoteS3,
         'gcp'   : RemoteGS,
+        'azure' : RemoteAzure,
         'ssh'   : RemoteSSH,
         'hdfs'  : RemoteHDFS,
         'local' : RemoteLOCAL,
diff --git a/dvc/remote/__init__.py b/dvc/remote/__init__.py
index 40a44b5d3..a6537c921 100644
--- a/dvc/remote/__init__.py
+++ b/dvc/remote/__init__.py
@@ -3,12 +3,13 @@ from dvc.remote.s3 import RemoteS3
 from dvc.remote.gs import RemoteGS
 from dvc.remote.hdfs import RemoteHDFS
 from dvc.remote.ssh import RemoteSSH
+from dvc.remote.azure import RemoteAzure
 
 from dvc.config import Config
 from dvc.exceptions import DvcException
 
 
-REMOTES = [RemoteHDFS, RemoteSSH, RemoteS3, RemoteGS, RemoteLOCAL]
+REMOTES = [RemoteHDFS, RemoteSSH, RemoteS3, RemoteGS, RemoteAzure, RemoteLOCAL]
 
 
 def supported_url(url):
diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
new file mode 100644
index 000000000..6f6b09f63
--- /dev/null
+++ b/dvc/remote/azure.py
@@ -0,0 +1,241 @@
+import os
+import re
+import time
+
+try:
+    from azure.storage.blob import BlockBlobService
+except ImportError:
+    BlockBlobService = None
+
+try:
+    from urlparse import urlparse
+except ImportError:
+    from urllib.parse import urlparse
+
+from dvc.logger import Logger
+from dvc.progress import progress
+from dvc.config import Config
+from dvc.remote.base import RemoteBase
+
+
+class Callback(object):
+    def __init__(self, name):
+        self.name = name
+
+    def __call__(self, current, total):
+        progress.update_target(self.name, current, total)
+
+
+class RemoteAzure(RemoteBase):
+    scheme = 'azure'
+    REGEX = (r'^azure://'
+             r'(ContainerName=(?P<container_name>[^;]+);)?'
+             r'(?P<connection_string>.+)?$')
+    REQUIRES = {'azure-storage-blob': BlockBlobService}
+    PARAM_ETAG = 'etag'
+    COPY_POLL_SECONDS = 5
+
+    def __init__(self, project, config):
+        super().__init__(project, config)
+        self.project = project
+
+        url = config.get(Config.SECTION_REMOTE_URL)
+        match = re.match(self.REGEX, url)
+
+        self.bucket = (
+            match.group('container_name')
+            or os.getenv('AZURE_STORAGE_CONTAINER_NAME'))
+
+        self.connection_string = (
+            match.group('connection_string')
+            or os.getenv('AZURE_STORAGE_CONNECTION_STRING'))
+
+        if not self.bucket:
+            raise ValueError('Azure Storage container name missing')
+
+        if not self.connection_string:
+            raise ValueError('Azure Storage connection string missing')
+
+        self.__blob_service = None
+
+    @property
+    def blob_service(self):
+        if self.__blob_service is None:
+            self.__blob_service = BlockBlobService(
+                connection_string=self.connection_string)
+            self.__blob_service.create_container(self.bucket)
+        return self.__blob_service
+
+    def _get_etag(self, bucket, key):
+        try:
+            blob = self.blob_service.get_blob_properties(bucket, key)
+            return blob.properties.etag
+        except Exception:
+            return None
+
+    def save_info(self, path_info):
+        if path_info['scheme'] != self.scheme:
+            raise NotImplementedError
+
+        return {self.PARAM_ETAG: self._get_etag(
+            path_info['bucket'], path_info['key'])}
+
+    def save(self, path_info):
+        if path_info['scheme'] != self.scheme:
+            raise NotImplementedError
+
+        etag = self._get_etag(path_info['bucket'], path_info['key'])
+        dest_key = '{}/{}'.format(etag[0:2], etag[2:])
+
+        self._copy_blob(to_bucket=self.bucket,
+                        to_key=dest_key,
+                        from_bucket=path_info['bucket'],
+                        from_key=path_info['key'])
+
+        return {self.PARAM_ETAG: etag}
+
+    def _copy_blob(self, to_bucket, to_key, from_bucket, from_key):
+        source = self.blob_service.make_blob_url(from_bucket, from_key)
+
+        copy = self.blob_service.copy_blob(to_bucket, to_key, source)
+
+        if self.COPY_POLL_SECONDS <= 0:
+            return
+
+        while copy.status != 'success':
+            time.sleep(self.COPY_POLL_SECONDS)
+            copy = self.blob_service.get_blob_properties(
+                to_bucket, to_key).properties.copy
+
+    def checkout(self, path_info, checksum_info):
+        if path_info['scheme'] != self.scheme:
+            raise NotImplementedError
+
+        etag = checksum_info.get(self.PARAM_ETAG, None)
+        if not etag:
+            return
+
+        self._copy_blob(to_bucket=path_info['bucket'],
+                        to_key=path_info['key'],
+                        from_bucket=self.bucket,
+                        from_key='{}/{}'.format(etag[0:2], etag[2:]))
+
+    def remove(self, path_info):
+        if path_info['scheme'] != self.scheme:
+            raise NotImplementedError
+
+        Logger.debug('Removing azure://{}/{}'.format(path_info['bucket'],
+                                                     path_info['key']))
+
+        self.blob_service.delete_blob(path_info['bucket'], path_info['key'])
+
+    def md5s_to_path_infos(self, md5s):
+        return [{
+            'scheme': self.scheme,
+            'bucket': self.bucket,
+            'key': '{}/{}'.format(md5[0:2], md5[2:])
+        } for md5 in md5s]
+
+    def exists(self, path_infos):
+        # TODO: why does S3 fetch all blobs first?
+        ret = []
+
+        for path_info in path_infos:
+            if path_info['scheme'] != self.scheme:
+                raise NotImplementedError
+
+            bucket = path_info['bucket']
+            key = path_info['key']
+
+            exists = self.blob_service.exists(bucket, key)
+            ret.append(exists)
+
+        return ret
+
+    def upload(self, paths, path_infos, names=None):
+        assert isinstance(paths, list)
+        assert isinstance(path_infos, list)
+        assert len(paths) == len(path_infos)
+        if not names:
+            names = len(paths) * [None]
+        else:
+            assert isinstance(names, list)
+            assert len(names) == len(paths)
+
+        for path, path_info, name in zip(paths, path_infos, names):
+            if path_info['scheme'] != self.scheme:
+                raise NotImplementedError
+
+            bucket = path_info['bucket']
+            key = path_info['key']
+
+            Logger.debug("Uploading '{}' to '{}/{}'".format(
+                path, bucket, key))
+
+            if not name:
+                name = os.path.basename(path)
+
+            cb = Callback(name)
+
+            try:
+                self.blob_service.create_blob_from_path(
+                    bucket, key, path, progress_callback=cb)
+            except Exception as ex:
+                Logger.error("Failed to upload '{}'".format(path), ex)
+            else:
+                progress.finish_target(name)
+
+    def download(self, path_infos, fnames, no_progress_bar=False, names=None):
+        assert isinstance(fnames, list)
+        assert isinstance(path_infos, list)
+        assert len(fnames) == len(path_infos)
+        if not names:
+            names = len(fnames) * [None]
+        else:
+            assert isinstance(names, list)
+            assert len(names) == len(fnames)
+
+        for fname, path_info, name in zip(fnames, path_infos, names):
+            if path_info['scheme'] != self.scheme:
+                raise NotImplementedError
+
+            bucket = path_info['bucket']
+            key = path_info['key']
+
+            Logger.debug("Downloading '{}/{}' to '{}'".format(
+                bucket, key, fname))
+
+            tmp_file = self.tmp_file(fname)
+            if not name:
+                name = os.path.basename(fname)
+
+            cb = None if no_progress_bar else Callback(name)
+
+            self._makedirs(fname)
+
+            try:
+                self.blob_service.get_blob_to_path(
+                    bucket, key, tmp_file, progress_callback=cb)
+            except Exception as exc:
+                Logger.error("Failed to download '{}/{}'".format(
+                    path_info['bucket'], path_info['key']), exc)
+                return
+
+            os.rename(tmp_file, fname)
+
+            if not no_progress_bar:
+                progress.finish_target(name)
+
+    def gc(self, checksum_infos):
+        used_etags = [info[self.PARAM_ETAG] for info in checksum_infos]
+
+        all_blobs = self.blob_service.list_blobs(self.bucket)
+
+        for blob in all_blobs:
+            etag = blob.properties.etag
+            if etag in used_etags:
+                continue
+            path_info = {'scheme': self.scheme,
+                         'key': blob.name,
+                         'bucket': self.bucket}
+            self.remove(path_info)
diff --git a/requirements.txt b/requirements.txt
index 938a5ef76..4fdf9385e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+azure-storage-blob==1.3.0
 boto3==1.7.4
 ply>=3.9 # See https://github.com/pyinstaller/pyinstaller/issues/1945
 configparser>=3.5.0
diff --git a/setup.py b/setup.py
index f6b2e6886..3f4a0291f 100644
--- a/setup.py
+++ b/setup.py
@@ -32,10 +32,13 @@ gs = [
 s3 = [
     "boto3==1.7.4",
 ]
+azure = [
+    "azure-storage-blob==1.3.0"
+]
 ssh = [
     "paramiko>=2.4.1",
 ]
-all_remotes = gs + s3 + ssh
+all_remotes = gs + s3 + azure + ssh
 
 setup(
     name='dvc',
@@ -51,6 +54,7 @@ setup(
         'all': all_remotes,
         'gs': gs,
         's3': s3,
+        'azure': azure,
         'ssh': ssh,
     },
     keywords='data science, data version control, machine learning',


Test patch:
diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py
index a134c754b..66f582d5e 100644
--- a/tests/test_data_cloud.py
+++ b/tests/test_data_cloud.py
@@ -9,7 +9,7 @@ import platform
 
 from dvc.main import main
 from dvc.config import Config, ConfigError
-from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH, RemoteHDFS
+from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteAzure, RemoteLOCAL, RemoteSSH, RemoteHDFS
 from dvc.remote.base import STATUS_OK, STATUS_NEW, STATUS_DELETED
 
 from tests.basic_env import TestDvc
@@ -53,6 +53,15 @@ def _should_test_gcp():
     return False
 
 
+def _should_test_azure():
+    if os.getenv("DVC_TEST_AZURE") == "true":
+        return True
+    elif os.getenv("DVC_TEST_AZURE") == "false":
+        return False
+
+    return os.getenv("AZURE_STORAGE_CONTAINER") and os.getenv("AZURE_STORAGE_CONNECTION_STRING")
+
+
 def _should_test_ssh():
     if os.getenv("DVC_TEST_SSH") == "true":
         return True
@@ -102,6 +111,10 @@ def get_gcp_url():
     return 'gs://' + get_gcp_storagepath()
 
 
+def get_azure_url():
+    return 'azure://'
+
+
 def sleep():
     # Sleep some time to simulate realistic behavior.
     # Some filesystems have a bad date resolution for
@@ -242,6 +255,17 @@ class TestRemoteGS(TestDataCloudBase):
         return RemoteGS
 
 
+class TestRemoteAzure(TestDataCloudBase):
+    def _should_test(self):
+        return _should_test_azure()
+
+    def _get_url(self):
+        return get_azure_url()
+
+    def _get_cloud_class(self):
+        return RemoteAzure
+
+
 class TestRemoteLOCAL(TestDataCloudBase):
     def _should_test(self):
         return True
@@ -433,6 +457,18 @@ class TestRemoteGSCLI(TestDataCloudCLIBase):
         self._test_cloud(TEST_REMOTE)
 
 
+class TestRemoteAzureCLI(TestDataCloudCLIBase):
+    def _should_test(self):
+        return _should_test_azure()
+
+    def _test(self):
+        url = get_azure_url()
+
+        self.main(['remote', 'add', TEST_REMOTE, url])
+
+        self._test_cloud(TEST_REMOTE)
+
+
 class TestDataCloudErrorCLI(TestDvc):
     def main_fail(self, args):
         ret = main(args)