From eaa2d559db7bc252e2cf04c5230abbf9da7f839d Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Wed, 12 Dec 2018 10:49:25 +0200
Subject: [PATCH] azure: support prefix

Fixes #1336

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/config.py            |  3 +++
 dvc/remote/azure.py      | 31 ++++++++++++++++++++-----------
 tests/test_data_cloud.py | 17 ++++++++++++++---
 3 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/dvc/config.py b/dvc/config.py
index 70a3b88e5..aaa34a879 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -150,6 +150,8 @@ class Config(object):
         SECTION_LOCAL_STORAGEPATH: str,
     }
 
+    SECTION_AZURE_CONNECTION_STRING = 'connection_string'
+
     SECTION_REMOTE_REGEX = r'^\s*remote\s*"(?P<name>.*)"\s*$'
     SECTION_REMOTE_FMT = 'remote "{}"'
     SECTION_REMOTE_URL = 'url'
@@ -177,6 +179,7 @@ class Config(object):
         Optional(SECTION_REMOTE_TIMEOUT): Use(int),
         Optional(SECTION_REMOTE_PASSWORD): str,
         Optional(SECTION_REMOTE_ASK_PASSWORD): And(str, is_bool, Use(to_bool)),
+        Optional(SECTION_AZURE_CONNECTION_STRING): str,
     }
 
     SECTION_STATE = 'state'
diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
index f971b4141..228c68551 100644
--- a/dvc/remote/azure.py
+++ b/dvc/remote/azure.py
@@ -8,6 +8,11 @@ try:
 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
@@ -24,9 +29,10 @@ class Callback(object):
 
 class RemoteAzure(RemoteBase):
     scheme = 'azure'
-    REGEX = (r'^azure://'
+    REGEX = (r'^azure://((?P<path>[^=;]*)|('
+             # backward compatibility
              r'(ContainerName=(?P<container_name>[^;]+);?)?'
-             r'(?P<connection_string>.+)?$')
+             r'(?P<connection_string>.+)?))$')
     REQUIRES = {'azure-storage-blob': BlockBlobService}
     PARAM_ETAG = 'etag'
     COPY_POLL_SECONDS = 5
@@ -36,17 +42,18 @@ class RemoteAzure(RemoteBase):
         self.project = project
 
         self.url = config.get(Config.SECTION_REMOTE_URL)
-        match = re.match(self.REGEX, self.url)
+        match = re.match(self.REGEX, self.url)  # backward compatibility
 
         self.bucket = (
-            match.group('container_name')
+            urlparse(self.url).netloc
+            or match.group('container_name')  # backward compatibility
             or os.getenv('AZURE_STORAGE_CONTAINER_NAME'))
 
-        # FIXME: currently Azure doesn't support prefixes
-        self.prefix = '/'
+        self.prefix = urlparse(self.url).path.lstrip('/') or '/'
 
         self.connection_string = (
-            match.group('connection_string')
+            config.get(Config.SECTION_AZURE_CONNECTION_STRING)
+            or match.group('connection_string')  # backward compatibility
             or os.getenv('AZURE_STORAGE_CONNECTION_STRING'))
 
         if not self.bucket:
@@ -138,7 +145,7 @@ class RemoteAzure(RemoteBase):
         return [{
             'scheme': self.scheme,
             'bucket': self.bucket,
-            'key': '{}/{}'.format(md5[0:2], md5[2:])
+            'key': '{}/{}/{}'.format(self.prefix, md5[0:2], md5[2:])
         } for md5 in md5s]
 
     def _all_keys(self, bucket, prefix):
@@ -235,8 +242,8 @@ class RemoteAzure(RemoteBase):
                     progress.finish_target(name)
 
     def _path_to_etag(self, path):
-        assert len(path.split('/')) == 2
-        return posixpath.dirname(path) + posixpath.basename(path)
+        relpath = posixpath.relpath(path, self.prefix)
+        return posixpath.dirname(relpath) + posixpath.basename(relpath)
 
     def _all(self):
         keys = self._all_keys(self.bucket, self.prefix)
@@ -253,7 +260,9 @@ class RemoteAzure(RemoteBase):
             if etag in used:
                 continue
             path_info = {'scheme': self.scheme,
-                         'key': posixpath.join(etag[0:2], etag[2:]),
+                         'key': posixpath.join(self.prefix,
+                                               etag[0:2],
+                                               etag[2:]),
                          'bucket': self.bucket}
             self.remove(path_info)
             removed = True
diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py
index 35149a324..bff0abb1c 100644
--- a/tests/test_data_cloud.py
+++ b/tests/test_data_cloud.py
@@ -141,10 +141,16 @@ def get_gcp_url():
     return 'gs://' + get_gcp_storagepath()
 
 
-def get_azure_url():
+def get_azure_url_compat():
     return 'azure://'
 
 
+def get_azure_url():
+    container_name = os.getenv("AZURE_STORAGE_CONTAINER_NAME")
+    assert container_name is not None
+    return 'azure://{}/{}'.format(container_name, str(uuid.uuid4()))
+
+
 def sleep():
     # Sleep some time to simulate realistic behavior.
     # Some filesystems have a bad date resolution for
@@ -312,17 +318,22 @@ class TestRemoteGS(TestDataCloudBase):
         return RemoteGS
 
 
-class TestRemoteAzure(TestDataCloudBase):
+class TestRemoteAzureCompat(TestDataCloudBase):
     def _should_test(self):
         return _should_test_azure()
 
     def _get_url(self):
-        return get_azure_url()
+        return get_azure_url_compat()
 
     def _get_cloud_class(self):
         return RemoteAzure
 
 
+class TestRemoteAzure(TestRemoteAzureCompat):
+    def _get_url(self):
+        return get_azure_url()
+
+
 class TestRemoteLOCAL(TestDataCloudBase):
     def _should_test(self):
         return True