diff --git a/.pylintrc b/.pylintrc index 6d81c3491..754000d2d 100644 --- a/.pylintrc +++ b/.pylintrc @@ -293,7 +293,7 @@ ignored-classes=optparse.Values,thread._local,_thread._local,Dvcfile # (useful for modules/projects where namespaces are manipulated during runtime # and thus existing member attributes cannot be deduced by static analysis. It # supports qualified module names, as well as Unix pattern matching. -ignored-modules= +ignored-modules=azure # Show a hint with possible names when a member name was not found. The aspect # of finding the hint is based on edit distance. diff --git a/dvc/config_schema.py b/dvc/config_schema.py index f8a023b96..df84458fe 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -179,7 +179,15 @@ SCHEMA = { "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 index 2271fb775..fddbc9395 100644 --- 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 - REQUIRES = { - "azure-storage-blob": "azure.storage.blob", - "knack": "knack", - } PARAM_CHECKSUM = "etag" DETAIL_FIELDS = frozenset(("etag", "size")) - - COPY_POLL_SECONDS = 5 - LIST_OBJECT_PAGE_SIZE = 5000 + REQUIRES = { + "adlfs": "adlfs", + "knack": "knack", + "azure-identity": "azure.identity", + } 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._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" + self.login_method, self.login_info = self._prepare_credentials(config) - self._credential = config.get("sas_token") or self._az_config.get( - "storage", "sas_token", None + def _prepare_credentials(self, config): + from azure.identity.aio import DefaultAzureCredential + + login_info = {} + login_info["connection_string"] = config.get( + "connection_string", + self._az_config.get("storage", "connection_string", None), ) - if not self._credential: - self._credential = self._az_config.get("storage", "key", 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) + ) + 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 @@ class AzureFileSystem(BaseFileSystem): @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 + def fs(self): + from adlfs import AzureBlobFileSystem + from azure.core.exceptions import AzureError - 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 - ) + 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 - logger.debug(f"Container name {self.bucket}") - container_client = blob_service.get_container_client(self.bucket) + return file_system - 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 + def _with_bucket(self, path): + if isinstance(path, self.PATH_CLS): + return f"{path.bucket}/{path.path}" + return path - return blob_service + def open( + self, path_info, mode="r", **kwargs + ): # pylint: disable=arguments-differ + return self.fs.open(self._with_bucket(path_info), mode=mode) - 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 exists(self, path_info, use_dvcignore=False): + return self.fs.exists(self._with_bucket(path_info)) - def _generate_download_url(self, path_info, expires=3600): - from azure.storage.blob import ( # pylint:disable=no-name-in-module - BlobSasPermissions, - generate_blob_sas, - ) + def _strip_bucket(self, entry): + _, entry = entry.split("/", 1) + return entry - 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 _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 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 - - 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 walk_files(self, path_info, **kwargs): + for file in self.ls(path_info, recursive=True): + yield path_info.replace(path=file) def remove(self, path_info): - if path_info.scheme != self.scheme: - raise NotImplementedError - - logger.debug(f"Removing {path_info}") - self.blob_service.get_blob_client( - path_info.bucket, path_info.path - ).delete_blob() + 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 index 4e63651f1..baf31f798 100644 --- a/setup.py +++ b/setup.py @@ -97,7 +97,7 @@ install_requires = [ gs = ["google-cloud-storage==1.19.0"] gdrive = ["pydrive2>=1.7.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"] diff --git a/tests/unit/remote/test_azure.py b/tests/unit/remote/test_azure.py index 397377cde..34bddd0d1 100644 --- 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):