minimal ADLS gen2 support

Author: pjbull

CONTRIBUTING.md

@@ -81,6 +81,15 @@ Finally, you may want to run your tests against live servers to ensure that the
 make test-live-cloud
 ```
 
+#### Azure live backend tests
+
+For Azure, you can test both against Azure Blob Storage backends and Azure Data Lake Storage Gen2 backends. To run these tests, you need to set connection strings for both of the backends by setting the following environment variables (in your `.env` file for local development). If `AZURE_STORAGE_GEN2_CONNECTION_STRING` is not set, only the blob storage backend will be tested.
+
+```bash
+AZURE_STORAGE_CONNECTION_STRING=your_connection_string
+AZURE_STORAGE_GEN2_CONNECTION_STRING=your_connection_string
+```
+
 You can copy `.env.example` to `.env` and fill in the credentials and bucket/container names for the providers you want to test against.  **Note that the live tests will create and delete files on the cloud provider.**
 
 You can also skip providers you do not have accounts for by commenting them out in the `rig` and `s3_like_rig` variables defined at the end of `tests/conftest.py`. 

HISTORY.md

@@ -8,6 +8,7 @@
 - fix: use native `exists()` method in `GSClient`. (PR [#420](https://github.com/drivendataorg/cloudpathlib/pull/420))
 - Enhancement: lazy instantiation of default client (PR [#432](https://github.com/drivendataorg/cloudpathlib/issues/432), Issue [#428](https://github.com/drivendataorg/cloudpathlib/issues/428))
 - Adds existence check before downloading in `download_to` (Issue [#430](https://github.com/drivendataorg/cloudpathlib/issues/430), PR [#432](https://github.com/drivendataorg/cloudpathlib/pull/432))
+- Implement minimal Azure Data Lake Storage Gen2 support (Issue [#161](https://github.com/drivendataorg/cloudpathlib/issues/161), PR [#450](https://github.com/drivendataorg/cloudpathlib/pull/450))
 
 ## v0.18.1 (2024-02-26)
 

cloudpathlib/azure/azblobclient.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timedelta
 import mimetypes
 import os
-from pathlib import Path, PurePosixPath
+from pathlib import Path
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
 
 
@@ -15,12 +15,15 @@
 try:
     from azure.core.exceptions import ResourceNotFoundError
     from azure.storage.blob import (
+        BlobPrefix,
         BlobSasPermissions,
         BlobServiceClient,
         BlobProperties,
         ContentSettings,
         generate_blob_sas,
     )
+
+    from azure.storage.filedatalake import DataLakeServiceClient, FileProperties
 except ModuleNotFoundError:
     implementation_registry["azure"].dependencies_loaded = False
 
@@ -39,6 +42,7 @@ def __init__(
         credential: Optional[Any] = None,
         connection_string: Optional[str] = None,
         blob_service_client: Optional["BlobServiceClient"] = None,
+        data_lake_client: Optional["DataLakeServiceClient"] = None,
         file_cache_mode: Optional[Union[str, FileCacheMode]] = None,
         local_cache_dir: Optional[Union[str, os.PathLike]] = None,
         content_type_method: Optional[Callable] = mimetypes.guess_type,
@@ -76,6 +80,10 @@ def __init__(
                 https://docs.microsoft.com/en-us/azure/storage/blobs/storage-quickstart-blobs-python#copy-your-credentials-from-the-azure-portal).
             blob_service_client (Optional[BlobServiceClient]): Instantiated [`BlobServiceClient`](
                 https://docs.microsoft.com/en-us/python/api/azure-storage-blob/azure.storage.blob.blobserviceclient?view=azure-python).
+            data_lake_client (Optional[DataLakeServiceClient]): Instantiated [`DataLakeServiceClient`](
+                https://learn.microsoft.com/en-us/python/api/azure-storage-file-datalake/azure.storage.filedatalake.datalakeserviceclient).
+                If None and `blob_service_client` is passed, we will create based on that.
+                Otherwise, will create based on passed credential, account_url, connection_string, or AZURE_STORAGE_CONNECTION_STRING env var
             file_cache_mode (Optional[Union[str, FileCacheMode]]): How often to clear the file cache; see
                 [the caching docs](https://cloudpathlib.drivendata.org/stable/caching/) for more information
                 about the options in cloudpathlib.eums.FileCacheMode.
@@ -94,27 +102,73 @@ def __init__(
         if connection_string is None:
             connection_string = os.getenv("AZURE_STORAGE_CONNECTION_STRING", None)
 
+        self.blob_service_client = None
+        self.data_lake_client = None
+
         if blob_service_client is not None:
             self.service_client = blob_service_client
+
+            # create from blob service client if not passed
+            if data_lake_client is None:
+                self.data_lake_client = DataLakeServiceClient(
+                    account_url=f"https://{self.service_client.account_name}.dfs.core.windows.net",
+                    credential=blob_service_client.credential,
+                )
+
+        if data_lake_client is not None:
+            self.data_lake_client = data_lake_client
+
         elif connection_string is not None:
             self.service_client = BlobServiceClient.from_connection_string(
                 conn_str=connection_string, credential=credential
             )
+            self.data_lake_client = DataLakeServiceClient.from_connection_string(
+                conn_str=connection_string, credential=credential
+            )
         elif account_url is not None:
             self.service_client = BlobServiceClient(account_url=account_url, credential=credential)
+            self.data_lake_client = DataLakeServiceClient(
+                account_url=account_url, credential=credential
+            )
         else:
             raise MissingCredentialsError(
                 "AzureBlobClient does not support anonymous instantiation. "
                 "Credentials are required; see docs for options."
             )
 
-    def _get_metadata(self, cloud_path: AzureBlobPath) -> Union["BlobProperties", Dict[str, Any]]:
-        blob = self.service_client.get_blob_client(
-            container=cloud_path.container, blob=cloud_path.blob
-        )
-        properties = blob.get_blob_properties()
+        self.hns_cache: Dict[str, bool] = {}
+
+    def _check_hns(self, cloud_path: AzureBlobPath) -> bool:
+        if cloud_path.container not in self.hns_cache:
+            hns_enabled: bool = self.service_client.get_account_information().get(
+                "is_hns_enabled", False
+            )  # type: ignore
+            self.hns_cache[cloud_path.container] = hns_enabled
+
+        return self.hns_cache[cloud_path.container]
+
+    def _get_metadata(
+        self, cloud_path: AzureBlobPath
+    ) -> Union["BlobProperties", "FileProperties", Dict[str, Any]]:
+        if self._check_hns(cloud_path):
+            # works on both files and directories
+            fsc = self.data_lake_client.get_file_system_client(cloud_path.container)  # type: ignore
+
+            if fsc is not None:
+                properties = fsc.get_file_client(cloud_path.blob).get_file_properties()
+
+            # no content settings on directory
+            properties["content_type"] = properties.get(
+                "content_settings", {"content_type": None}
+            ).get("content_type")
 
-        properties["content_type"] = properties.content_settings.content_type
+        else:
+            blob = self.service_client.get_blob_client(
+                container=cloud_path.container, blob=cloud_path.blob
+            )
+            properties = blob.get_blob_properties()
+
+            properties["content_type"] = properties.content_settings.content_type
 
         return properties
 
@@ -155,8 +209,17 @@ def _is_file_or_dir(self, cloud_path: AzureBlobPath) -> Optional[str]:
             return "dir"
 
         try:
-            self._get_metadata(cloud_path)
-            return "file"
+            meta = self._get_metadata(cloud_path)
+
+            # if hns, has is_directory property; else if not hns, _get_metadata will raise if not a file
+            return (
+                "dir"
+                if meta.get("is_directory", False)
+                or meta.get("metadata", {}).get("hdi_isfolder", False)
+                else "file"
+            )
+
+        # thrown if not HNS and file does not exist _or_ is dir; check if is dir instead
         except ResourceNotFoundError:
             prefix = cloud_path.blob
             if prefix and not prefix.endswith("/"):
@@ -181,17 +244,14 @@ def _exists(self, cloud_path: AzureBlobPath) -> bool:
     def _list_dir(
         self, cloud_path: AzureBlobPath, recursive: bool = False
     ) -> Iterable[Tuple[AzureBlobPath, bool]]:
-        # shortcut if listing all available containers
         if not cloud_path.container:
-            if recursive:
-                raise NotImplementedError(
-                    "Cannot recursively list all containers and contents; you can get all the containers then recursively list each separately."
-                )
+            for container in self.service_client.list_containers():
+                yield self.CloudPath(f"az://{container.name}"), True
 
-            yield from (
-                (self.CloudPath(f"az://{c.name}"), True)
-                for c in self.service_client.list_containers()
-            )
+                if not recursive:
+                    continue
+
+                yield from self._list_dir(self.CloudPath(f"az://{container.name}"), recursive=True)
             return
 
         container_client = self.service_client.get_container_client(cloud_path.container)
@@ -200,30 +260,24 @@ def _list_dir(
         if prefix and not prefix.endswith("/"):
             prefix += "/"
 
-        yielded_dirs = set()
-
-        # NOTE: Not recursive may be slower than necessary since it just filters
-        #   the recursive implementation
-        for o in container_client.list_blobs(name_starts_with=prefix):
-            # get directory from this path
-            for parent in PurePosixPath(o.name[len(prefix) :]).parents:
-                # if we haven't surfaced this directory already
-                if parent not in yielded_dirs and str(parent) != ".":
-                    # skip if not recursive and this is beyond our depth
-                    if not recursive and "/" in str(parent):
-                        continue
+        if self._check_hns(cloud_path):
+            file_system_client = self.data_lake_client.get_file_system_client(cloud_path.container)  # type: ignore
+            paths = file_system_client.get_paths(path=cloud_path.blob, recursive=recursive)
 
-                    yield (
-                        self.CloudPath(f"az://{cloud_path.container}/{prefix}{parent}"),
-                        True,  # is a directory
-                    )
-                    yielded_dirs.add(parent)
+            for path in paths:
+                yield self.CloudPath(f"az://{cloud_path.container}/{path.name}"), path.is_directory
 
-            # skip file if not recursive and this is beyond our depth
-            if not recursive and "/" in o.name[len(prefix) :]:
-                continue
+        else:
+            if not recursive:
+                blobs = container_client.walk_blobs(name_starts_with=prefix)
+            else:
+                blobs = container_client.list_blobs(name_starts_with=prefix)
 
-            yield (self.CloudPath(f"az://{cloud_path.container}/{o.name}"), False)  # is a file
+            for blob in blobs:
+                # walk_blobs returns folders with a trailing slash
+                blob_path = blob.name.rstrip("/")
+                blob_cloud_path = self.CloudPath(f"az://{cloud_path.container}/{blob_path}")
+                yield blob_cloud_path, isinstance(blob, BlobPrefix)
 
     def _move_file(
         self, src: AzureBlobPath, dst: AzureBlobPath, remove_src: bool = True
@@ -253,6 +307,10 @@ def _move_file(
     def _remove(self, cloud_path: AzureBlobPath, missing_ok: bool = True) -> None:
         file_or_dir = self._is_file_or_dir(cloud_path)
         if file_or_dir == "dir":
+            if self._check_hns(cloud_path):
+                _hns_rmtree(self.data_lake_client, cloud_path.container, cloud_path.blob)
+                return
+
             blobs = [
                 b.blob for b, is_dir in self._list_dir(cloud_path, recursive=True) if not is_dir
             ]
@@ -313,4 +371,15 @@ def _generate_presigned_url(
         return url
 
 
+def _hns_rmtree(data_lake_client, container, directory):
+    """Stateless implementation so can be used in test suite cleanup as well.
+
+    If hierarchical namespace is enabled, delete the directory and all its contents.
+    (The non-HNS version is implemented in `_remove`, but will leave empty folders in HNS).
+    """
+    file_system_client = data_lake_client.get_file_system_client(container)
+    directory_client = file_system_client.get_directory_client(directory)
+    directory_client.delete_directory()
+
+
 AzureBlobClient.AzureBlobPath = AzureBlobClient.CloudPath  # type: ignore

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-azure = ["azure-storage-blob>=12"]
+azure = ["azure-storage-blob>=12", "azure-storage-file-datalake>=12"]
 gs = ["google-cloud-storage"]
 s3 = ["boto3>=1.34.0"]
 all = ["cloudpathlib[azure]", "cloudpathlib[gs]", "cloudpathlib[s3]"]