Repository: iterative/dvc

Issue:
azure: QA our implementation
People have been complaining about our azure remotes having other issues with authentication, e.g.

https://github.com/iterative/dvc/issues/4534
https://github.com/iterative/dvc/issues/4635

Azure support is very important for us, but unfortunately, no one in the core team is using it regularly (except for the tests that we have, which are clearly not enough). Need to take a look into it, check different auth options, check performance, etc.


Discussion:

Resolving patch:
diff --git a/dvc/config_schema.py b/dvc/config_schema.py
--- a/dvc/config_schema.py
+++ b/dvc/config_schema.py
@@ -179,7 +179,15 @@ class RelPath(str):
                     "webhdfs_alias": str,
                     **REMOTE_COMMON,
                 },
-                "azure": {"connection_string": str, **REMOTE_COMMON},
+                "azure": {
+                    "connection_string": str,
+                    "sas_token": str,
+                    "account_name": str,
+                    "tenant_id": str,
+                    "client_id": str,
+                    "client_secret": str,
+                    **REMOTE_COMMON,
+                },
                 "oss": {
                     "oss_key_id": str,
                     "oss_key_secret": str,
diff --git a/dvc/fs/azure.py b/dvc/fs/azure.py
--- a/dvc/fs/azure.py
+++ b/dvc/fs/azure.py
@@ -1,58 +1,94 @@
 import logging
 import os
+import shutil
 import threading
-from datetime import datetime, timedelta
 
 from funcy import cached_property, wrap_prop
 
+from dvc.exceptions import DvcException
 from dvc.path_info import CloudURLInfo
 from dvc.progress import Tqdm
 from dvc.scheme import Schemes
+from dvc.utils import format_link
 
 from .base import BaseFileSystem
 
 logger = logging.getLogger(__name__)
+_DEFAULT_CREDS_STEPS = (
+    "https://azuresdkdocs.blob.core.windows.net/$web/python/"
+    "azure-identity/1.4.0/azure.identity.html#azure.identity"
+    ".DefaultAzureCredential"
+)
+
+
+class AzureAuthError(DvcException):
+    pass
 
 
 class AzureFileSystem(BaseFileSystem):
     scheme = Schemes.AZURE
     PATH_CLS = CloudURLInfo
+    PARAM_CHECKSUM = "etag"
+    DETAIL_FIELDS = frozenset(("etag", "size"))
     REQUIRES = {
-        "azure-storage-blob": "azure.storage.blob",
+        "adlfs": "adlfs",
         "knack": "knack",
+        "azure-identity": "azure.identity",
     }
-    PARAM_CHECKSUM = "etag"
-    DETAIL_FIELDS = frozenset(("etag", "size"))
-
-    COPY_POLL_SECONDS = 5
-    LIST_OBJECT_PAGE_SIZE = 5000
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
 
-        url = config.get("url", "azure://")
+        url = config.get("url")
         self.path_info = self.PATH_CLS(url)
         self.bucket = self.path_info.bucket
 
         if not self.bucket:
             container = self._az_config.get("storage", "container_name", None)
-            self.path_info = self.PATH_CLS(f"azure://{container}")
-            self.bucket = self.path_info.bucket
+            url = f"azure://{container}"
 
-        self._conn_str = config.get(
-            "connection_string"
-        ) or self._az_config.get("storage", "connection_string", None)
+        self.path_info = self.PATH_CLS(url)
+        self.bucket = self.path_info.bucket
+
+        self.login_method, self.login_info = self._prepare_credentials(config)
 
-        self._account_url = None
-        if not self._conn_str:
-            name = self._az_config.get("storage", "account", None)
-            self._account_url = f"https://{name}.blob.core.windows.net"
+    def _prepare_credentials(self, config):
+        from azure.identity.aio import DefaultAzureCredential
 
-        self._credential = config.get("sas_token") or self._az_config.get(
-            "storage", "sas_token", None
+        login_info = {}
+        login_info["connection_string"] = config.get(
+            "connection_string",
+            self._az_config.get("storage", "connection_string", None),
+        )
+        login_info["account_name"] = config.get(
+            "account_name", self._az_config.get("storage", "account", None)
+        )
+        login_info["account_key"] = config.get(
+            "account_key", self._az_config.get("storage", "key", None)
+        )
+        login_info["sas_token"] = config.get(
+            "sas_token", self._az_config.get("storage", "sas_token", None)
         )
-        if not self._credential:
-            self._credential = self._az_config.get("storage", "key", None)
+        login_info["tenant_id"] = config.get("tenant_id")
+        login_info["client_id"] = config.get("client_id")
+        login_info["client_secret"] = config.get("client_secret")
+
+        if not any(login_info.values()):
+            login_info["credentials"] = DefaultAzureCredential()
+
+        for login_method, required_keys in [  # noqa
+            ("connection string", ["connection_string"]),
+            ("account key", ["account_name", "account_key"]),
+            ("SAS token", ["account_name", "sas_token"]),
+            ("anonymous login", ["account_name"]),
+            (f"default credentials ({_DEFAULT_CREDS_STEPS})", ["credentials"]),
+        ]:
+            if all(login_info.get(key) is not None for key in required_keys):
+                break
+        else:
+            login_method = None
+
+        return login_method, login_info
 
     @cached_property
     def _az_config(self):
@@ -68,158 +104,109 @@ def _az_config(self):
 
     @wrap_prop(threading.Lock())
     @cached_property
-    def blob_service(self):
-        # pylint: disable=no-name-in-module
-        from azure.core.exceptions import (
-            HttpResponseError,
-            ResourceNotFoundError,
-        )
-        from azure.storage.blob import BlobServiceClient
-
-        if self._conn_str:
-            logger.debug(f"Using connection string '{self._conn_str}'")
-            blob_service = BlobServiceClient.from_connection_string(
-                self._conn_str, credential=self._credential
-            )
-        else:
-            logger.debug(f"Using account url '{self._account_url}'")
-            blob_service = BlobServiceClient(
-                self._account_url, credential=self._credential
-            )
-
-        logger.debug(f"Container name {self.bucket}")
-        container_client = blob_service.get_container_client(self.bucket)
-
-        try:  # verify that container exists
-            container_client.get_container_properties()
-        except ResourceNotFoundError:
-            container_client.create_container()
-        except HttpResponseError as exc:
-            # client may not have account-level privileges
-            if exc.status_code != 403:
-                raise
-
-        return blob_service
-
-    def get_etag(self, path_info):
-        blob_client = self.blob_service.get_blob_client(
-            path_info.bucket, path_info.path
-        )
-        etag = blob_client.get_blob_properties().etag
-        return etag.strip('"')
-
-    def _generate_download_url(self, path_info, expires=3600):
-        from azure.storage.blob import (  # pylint:disable=no-name-in-module
-            BlobSasPermissions,
-            generate_blob_sas,
-        )
-
-        expires_at = datetime.utcnow() + timedelta(seconds=expires)
-
-        blob_client = self.blob_service.get_blob_client(
-            path_info.bucket, path_info.path
-        )
-
-        sas_token = generate_blob_sas(
-            blob_client.account_name,
-            blob_client.container_name,
-            blob_client.blob_name,
-            account_key=blob_client.credential.account_key,
-            permission=BlobSasPermissions(read=True),
-            expiry=expires_at,
-        )
-        return blob_client.url + "?" + sas_token
-
-    def exists(self, path_info, use_dvcignore=True):
-        paths = self._list_paths(path_info.bucket, path_info.path)
-        return any(path_info.path == path for path in paths)
+    def fs(self):
+        from adlfs import AzureBlobFileSystem
+        from azure.core.exceptions import AzureError
+
+        try:
+            file_system = AzureBlobFileSystem(**self.login_info)
+            if self.bucket not in [
+                container.rstrip("/") for container in file_system.ls("/")
+            ]:
+                file_system.mkdir(self.bucket)
+        except (ValueError, AzureError) as e:
+            raise AzureAuthError(
+                f"Authentication to Azure Blob Storage via {self.login_method}"
+                " failed.\nLearn more about configuration settings at"
+                f" {format_link('https://man.dvc.org/remote/modify')}"
+            ) from e
+
+        return file_system
+
+    def _with_bucket(self, path):
+        if isinstance(path, self.PATH_CLS):
+            return f"{path.bucket}/{path.path}"
+        return path
+
+    def open(
+        self, path_info, mode="r", **kwargs
+    ):  # pylint: disable=arguments-differ
+        return self.fs.open(self._with_bucket(path_info), mode=mode)
 
-    def _list_paths(self, bucket, prefix):
-        container_client = self.blob_service.get_container_client(bucket)
-        for blob in container_client.list_blobs(name_starts_with=prefix):
-            yield blob.name
+    def exists(self, path_info, use_dvcignore=False):
+        return self.fs.exists(self._with_bucket(path_info))
 
-    def walk_files(self, path_info, **kwargs):
-        if not kwargs.pop("prefix", False):
-            path_info = path_info / ""
-        for fname in self._list_paths(
-            path_info.bucket, path_info.path, **kwargs
-        ):
-            if fname.endswith("/"):
-                continue
+    def _strip_bucket(self, entry):
+        _, entry = entry.split("/", 1)
+        return entry
 
-            yield path_info.replace(path=fname)
+    def _strip_buckets(self, entries, detail, prefix=None):
+        for entry in entries:
+            if detail:
+                entry = entry.copy()
+                entry["name"] = self._strip_bucket(entry["name"])
+            else:
+                entry = self._strip_bucket(
+                    f"{prefix}/{entry}" if prefix else entry
+                )
+            yield entry
 
     def ls(
         self, path_info, detail=False, recursive=False
     ):  # pylint: disable=arguments-differ
-        assert recursive
+        path = self._with_bucket(path_info)
+        if recursive:
+            for root, _, files in self.fs.walk(path, detail=detail):
+                if detail:
+                    files = files.values()
+                yield from self._strip_buckets(files, detail, prefix=root)
+            return None
 
-        container_client = self.blob_service.get_container_client(
-            path_info.bucket
-        )
-        for blob in container_client.list_blobs(
-            name_starts_with=path_info.path
-        ):
-            if detail:
-                yield {
-                    "type": "file",
-                    "name": blob.name,
-                    "size": blob.size,
-                    "etag": blob.etag.strip('"'),
-                }
-            else:
-                yield blob.name
+        yield from self._strip_buckets(self.ls(path, detail=detail), detail)
 
-    def remove(self, path_info):
-        if path_info.scheme != self.scheme:
-            raise NotImplementedError
+    def walk_files(self, path_info, **kwargs):
+        for file in self.ls(path_info, recursive=True):
+            yield path_info.replace(path=file)
 
-        logger.debug(f"Removing {path_info}")
-        self.blob_service.get_blob_client(
-            path_info.bucket, path_info.path
-        ).delete_blob()
+    def remove(self, path_info):
+        self.fs.delete(self._with_bucket(path_info))
 
     def info(self, path_info):
-        blob_client = self.blob_service.get_blob_client(
-            path_info.bucket, path_info.path
-        )
-        properties = blob_client.get_blob_properties()
-        return {
-            "type": "file",
-            "size": properties.size,
-            "etag": properties.etag.strip('"'),
-        }
+        info = self.fs.info(self._with_bucket(path_info)).copy()
+        info["name"] = self._strip_bucket(info["name"])
+        return info
 
     def _upload_fobj(self, fobj, to_info):
-        blob_client = self.blob_service.get_blob_client(
-            to_info.bucket, to_info.path
-        )
-        blob_client.upload_blob(fobj, overwrite=True)
+        with self.open(to_info, "wb") as fdest:
+            shutil.copyfileobj(fobj, fdest)
 
     def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+        self, from_file, to_info, name=None, no_progress_bar=False, **kwargs
     ):
         total = os.path.getsize(from_file)
         with open(from_file, "rb") as fobj:
             self.upload_fobj(
                 fobj,
-                to_info,
+                self._with_bucket(to_info),
                 desc=name,
                 total=total,
                 no_progress_bar=no_progress_bar,
             )
+        self.fs.invalidate_cache(self._with_bucket(to_info.parent))
 
     def _download(
-        self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
+        self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
     ):
-        blob_client = self.blob_service.get_blob_client(
-            from_info.bucket, from_info.path
-        )
-        total = blob_client.get_blob_properties().size
-        stream = blob_client.download_blob()
-        with open(to_file, "wb") as fobj:
+        total = self.getsize(self._with_bucket(from_info))
+        with self.open(from_info, "rb") as fobj:
             with Tqdm.wrapattr(
-                fobj, "write", desc=name, total=total, disable=no_progress_bar
+                fobj,
+                "read",
+                desc=name,
+                disable=no_progress_bar,
+                bytes=True,
+                total=total,
+                **pbar_args,
             ) as wrapped:
-                stream.readinto(wrapped)
+                with open(to_file, "wb") as fdest:
+                    shutil.copyfileobj(wrapped, fdest)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -96,7 +96,7 @@ def run(self):
 gs = ["google-cloud-storage==1.19.0"]
 gdrive = ["pydrive2>=1.6.3", "six >= 1.13.0"]
 s3 = ["boto3>=1.9.201"]
-azure = ["azure-storage-blob>=12.0", "knack"]
+azure = ["adlfs>=0.6.3", "azure.identity>=1.4.0", "knack"]
 # https://github.com/Legrandin/pycryptodome/issues/465
 oss = ["oss2==2.6.1", "pycryptodome<3.9.9"]
 ssh = ["paramiko[invoke]>=2.7.0"]


Test patch:
diff --git a/tests/unit/remote/test_azure.py b/tests/unit/remote/test_azure.py
--- a/tests/unit/remote/test_azure.py
+++ b/tests/unit/remote/test_azure.py
@@ -17,7 +17,6 @@ def test_init_env_var(monkeypatch, dvc):
     config = {"url": "azure://"}
     fs = AzureFileSystem(dvc, config)
     assert fs.path_info == "azure://" + container_name
-    assert fs._conn_str == connection_string
 
 
 def test_init(dvc):
@@ -26,7 +25,6 @@ def test_init(dvc):
     config = {"url": url, "connection_string": connection_string}
     fs = AzureFileSystem(dvc, config)
     assert fs.path_info == url
-    assert fs._conn_str == connection_string
 
 
 def test_info(tmp_dir, azure):