Support for WebHDFS (#4747)

* Add support for WebHDFS

* Style and linting fixes

* Implement abstract methods and disable no-member pylint errors

* Add init unit test

* Yield client and fix redefinition of hdfs

* Get dynamic hdfs Docker port

* Add support for open in WebHDFSTree

* Remove hdfscli_config fallbacks in WebHDFSTree

* Use port 50070 for webhdfs

* WebHDFSTree supports http URL and proper walking

* Fix failing unit test after refactoring

* Add overwrite argument to upload

* Add remaining func tests for webhdfs

* Remove print statements

* Allow dynamic attributes in pylint for hdfs library

* Comment error on instantiating HDFS client from config

* Remove depth argument from call to walk

* Update dvc/tree/webhdfs.py

* Remove unnecessary deque

* Check exception messages hdfs client failure from config

* Return client directly

Co-authored-by: Ruslan Kuprieiev <[email protected]>

* Split long line

* Add size to HashInfo

* Implement move via hdfs.rename

* Add missing line

* Make parent directories before moving file in WebHDFS

* Add total to progress bar

Co-authored-by: Ruslan Kuprieiev <[email protected]>

.pylintrc

@@ -265,7 +265,7 @@ contextmanager-decorators=contextlib.contextmanager
 # List of members which are set dynamically and missed by pylint inference
 # system, and so shouldn't trigger E1101 when accessed. Python regular
 # expressions are accepted.
-generated-members=pytest.lazy_fixture,logging.TRACE,logger.trace,sys.getwindowsversion
+generated-members=hdfs.*,pytest.lazy_fixture,logging.TRACE,logger.trace,sys.getwindowsversion
 
 # Tells whether missing members accessed in mixin class should be ignored. A
 # mixin class is detected if its name ends with "mixin" (case insensitive).

dvc/cache/__init__.py

@@ -36,7 +36,13 @@ class Cache:
     """
 
     CACHE_DIR = "cache"
-    CLOUD_SCHEMES = [Schemes.S3, Schemes.GS, Schemes.SSH, Schemes.HDFS]
+    CLOUD_SCHEMES = [
+        Schemes.S3,
+        Schemes.GS,
+        Schemes.SSH,
+        Schemes.HDFS,
+        Schemes.WEBHDFS,
+    ]
 
     def __init__(self, repo):
         self.repo = repo

dvc/config.py

@@ -149,6 +149,7 @@ SCHEMA = {
         "s3": str,
         "gs": str,
         "hdfs": str,
+        "webhdfs": str,
         "ssh": str,
         "azure": str,
         # This is for default local cache
@@ -195,6 +196,13 @@ SCHEMA = {
                     **REMOTE_COMMON,
                 },
                 "hdfs": {"user": str, **REMOTE_COMMON},
+                "webhdfs": {
+                    "hdfscli_config": str,
+                    "webhdfs_token": str,
+                    "user": str,
+                    "webhdfs_alias": str,
+                    **REMOTE_COMMON,
+                },
                 "azure": {"connection_string": str, **REMOTE_COMMON},
                 "oss": {
                     "oss_key_id": str,

dvc/dependency/__init__.py

@@ -13,6 +13,7 @@ from dvc.dependency.s3 import S3Dependency
 from dvc.dependency.ssh import SSHDependency
 from dvc.dependency.webdav import WebDAVDependency
 from dvc.dependency.webdavs import WebDAVSDependency
+from dvc.dependency.webhdfs import WebHDFSDependency
 from dvc.output.base import BaseOutput
 from dvc.scheme import Schemes
 
@@ -28,7 +29,8 @@ DEPS = [
     S3Dependency,
     SSHDependency,
     WebDAVDependency,
-    WebDAVSDependency
+    WebDAVSDependency,
+    WebHDFSDependency,
     # NOTE: LocalDependency is the default choice
 ]
 
@@ -43,6 +45,7 @@ DEP_MAP = {
     Schemes.HTTPS: HTTPSDependency,
     Schemes.WEBDAV: WebDAVDependency,
     Schemes.WEBDAVS: WebDAVSDependency,
+    Schemes.WEBHDFS: WebHDFSDependency,
 }
 
 

dvc/dependency/webhdfs.py

@@ -0,0 +1,6 @@
+from dvc.dependency.base import BaseDependency
+from dvc.output.webhdfs import WebHDFSOutput
+
+
+class WebHDFSDependency(BaseDependency, WebHDFSOutput):
+    pass

dvc/output/__init__.py

@@ -11,18 +11,21 @@ from dvc.output.hdfs import HDFSOutput
 from dvc.output.local import LocalOutput
 from dvc.output.s3 import S3Output
 from dvc.output.ssh import SSHOutput
+from dvc.output.webhdfs import WebHDFSOutput
 from dvc.scheme import Schemes
 
 from ..tree import get_cloud_tree
 from ..tree.hdfs import HDFSTree
 from ..tree.local import LocalTree
 from ..tree.s3 import S3Tree
+from ..tree.webhdfs import WebHDFSTree
 
 OUTS = [
     HDFSOutput,
     S3Output,
     GSOutput,
     SSHOutput,
+    WebHDFSOutput,
     # NOTE: LocalOutput is the default choice
 ]
 
@@ -32,6 +35,7 @@ OUTS_MAP = {
     Schemes.GS: GSOutput,
     Schemes.SSH: SSHOutput,
     Schemes.LOCAL: LocalOutput,
+    Schemes.WEBHDFS: WebHDFSOutput,
 }
 
 CHECKSUM_SCHEMA = Any(
@@ -52,6 +56,7 @@ CHECKSUMS_SCHEMA = {
     LocalTree.PARAM_CHECKSUM: CHECKSUM_SCHEMA,
     S3Tree.PARAM_CHECKSUM: CHECKSUM_SCHEMA,
     HDFSTree.PARAM_CHECKSUM: CHECKSUM_SCHEMA,
+    WebHDFSTree.PARAM_CHECKSUM: CHECKSUM_SCHEMA,
 }
 
 SCHEMA = CHECKSUMS_SCHEMA.copy()

dvc/output/webhdfs.py

@@ -0,0 +1,7 @@
+from dvc.output.base import BaseOutput
+
+from ..tree.webhdfs import WebHDFSTree
+
+
+class WebHDFSOutput(BaseOutput):
+    TREE_CLS = WebHDFSTree

dvc/scheme.py

@@ -1,6 +1,7 @@
 class Schemes:
     SSH = "ssh"
     HDFS = "hdfs"
+    WEBHDFS = "webhdfs"
     S3 = "s3"
     AZURE = "azure"
     HTTP = "http"

dvc/tree/__init__.py

@@ -13,6 +13,7 @@ from .s3 import S3Tree
 from .ssh import SSHTree
 from .webdav import WebDAVTree
 from .webdavs import WebDAVSTree
+from .webhdfs import WebHDFSTree
 
 TREES = [
     AzureTree,
@@ -26,6 +27,7 @@ TREES = [
     OSSTree,
     WebDAVTree,
     WebDAVSTree,
+    WebHDFSTree
     # NOTE: LocalTree is the default
 ]
 

dvc/tree/webhdfs.py

@@ -0,0 +1,164 @@
+import logging
+import os
+import threading
+from contextlib import contextmanager
+from urllib.parse import urlparse
+
+from funcy import cached_property, wrap_prop
+
+from dvc.hash_info import HashInfo
+from dvc.path_info import CloudURLInfo
+from dvc.progress import Tqdm
+from dvc.scheme import Schemes
+
+from .base import BaseTree
+
+logger = logging.getLogger(__name__)
+
+
+def update_pbar(pbar, total):
+    """Update pbar to accept the two arguments passed by hdfs"""
+
+    def update(_, bytes_transfered):
+        if bytes_transfered == -1:
+            pbar.update_to(total)
+            return
+        pbar.update_to(bytes_transfered)
+
+    return update
+
+
+class WebHDFSTree(BaseTree):
+    scheme = Schemes.WEBHDFS
+    PATH_CLS = CloudURLInfo
+    REQUIRES = {"hdfs": "hdfs"}
+    PARAM_CHECKSUM = "checksum"
+
+    def __init__(self, repo, config):
+        super().__init__(repo, config)
+
+        self.path_info = None
+        url = config.get("url")
+        if not url:
+            return
+
+        parsed = urlparse(url)
+        user = parsed.username or config.get("user")
+
+        self.path_info = self.PATH_CLS.from_parts(
+            scheme="webhdfs",
+            host=parsed.hostname,
+            user=user,
+            port=parsed.port,
+            path=parsed.path,
+        )
+
+        self.hdfscli_config = config.get("hdfscli_config")
+        self.token = config.get("webhdfs_token")
+        self.alias = config.get("webhdfs_alias")
+
+    @wrap_prop(threading.Lock())
+    @cached_property
+    def hdfs_client(self):
+        import hdfs
+
+        logger.debug("URL: %s", self.path_info)
+        logger.debug("HDFSConfig: %s", self.hdfscli_config)
+
+        try:
+            return hdfs.config.Config(self.hdfscli_config).get_client(
+                self.alias
+            )
+        except hdfs.util.HdfsError as exc:
+            exc_msg = str(exc)
+            errors = (
+                "No alias specified",
+                "Invalid configuration file",
+                f"Alias {self.alias} not found",
+            )
+            if not any(err in exc_msg for err in errors):
+                raise
+
+            http_url = f"http://{self.path_info.host}:{self.path_info.port}"
+
+            if self.token is not None:
+                client = hdfs.TokenClient(http_url, token=self.token, root="/")
+            else:
+                client = hdfs.InsecureClient(
+                    http_url, user=self.path_info.user, root="/"
+                )
+
+        return client
+
+    @contextmanager
+    def open(self, path_info, mode="r", encoding=None):
+        assert mode in {"r", "rt", "rb"}
+
+        with self.hdfs_client.read(
+            path_info.path, encoding=encoding
+        ) as reader:
+            yield reader.read()
+
+    def walk_files(self, path_info, **kwargs):
+        if not self.exists(path_info):
+            return
+
+        root = path_info.path
+        for path, _, files in self.hdfs_client.walk(root):
+            for file_ in files:
+                path = os.path.join(path, file_)
+                yield path_info.replace(path=path)
+
+    def remove(self, path_info):
+        if path_info.scheme != self.scheme:
+            raise NotImplementedError
+
+        self.hdfs_client.delete(path_info.path)
+
+    def exists(self, path_info, use_dvcignore=True):
+        assert not isinstance(path_info, list)
+        assert path_info.scheme == "webhdfs"
+
+        status = self.hdfs_client.status(path_info.path, strict=False)
+        return status is not None
+
+    def get_file_hash(self, path_info):
+        checksum = self.hdfs_client.checksum(path_info.path)
+        hash_info = HashInfo(self.PARAM_CHECKSUM, checksum["bytes"])
+
+        hash_info.size = self.hdfs_client.status(path_info.path)["length"]
+        return hash_info
+
+    def copy(self, from_info, to_info, **_kwargs):
+        with self.hdfs_client.read(from_info.path) as reader:
+            content = reader.read()
+        self.hdfs_client.write(to_info.path, data=content)
+
+    def move(self, from_info, to_info, mode=None):
+        self.hdfs_client.makedirs(to_info.parent.path)
+        self.hdfs_client.rename(from_info.path, to_info.path)
+
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
+        total = os.path.getsize(from_file)
+        with Tqdm(
+            desc=name, total=total, disable=no_progress_bar, bytes=True
+        ) as pbar:
+            self.hdfs_client.upload(
+                to_info.path,
+                from_file,
+                overwrite=True,
+                progress=update_pbar(pbar, total),
+            )
+
+    def _download(
+        self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
+    ):
+        total = self.hdfs_client.status(from_info.path)["length"]
+        with Tqdm(
+            desc=name, total=total, disable=no_progress_bar, bytes=True
+        ) as pbar:
+            self.hdfs_client.download(
+                from_info.path, to_file, progress=update_pbar(pbar, total)
+            )

setup.py

@@ -98,13 +98,14 @@ ssh = ["paramiko[invoke]>=2.7.0"]
 
 # Remove the env marker if/when pyarrow is available for Python3.9
 hdfs = ["pyarrow>=2.0.0;  python_version < '3.9'"]
+webhdfs = ["hdfs==2.5.8"]
 webdav = ["webdavclient3>=3.14.5"]
 # gssapi should not be included in all_remotes, because it doesn't have wheels
 # for linux and mac, so it will fail to compile if user doesn't have all the
 # requirements, including kerberos itself. Once all the wheels are available,
 # we can start shipping it by default.
 ssh_gssapi = ["paramiko[invoke,gssapi]>=2.7.0"]
-all_remotes = gs + s3 + azure + ssh + oss + gdrive + hdfs + webdav
+all_remotes = gs + s3 + azure + ssh + oss + gdrive + hdfs + webhdfs + webdav
 
 # Extra dependecies to run tests
 tests_requirements = [
@@ -173,6 +174,7 @@ setup(
         "ssh": ssh,
         "ssh_gssapi": ssh_gssapi,
         "hdfs": hdfs,
+        "webhdfs": webhdfs,
         "webdav": webdav,
         "tests": tests_requirements,
     },

tests/docker-compose.yml

@@ -11,9 +11,12 @@ services:
       - "8880"
   hdfs:
     image: rkuprieiev/docker-hdfs
+    hostname: localhost
     ports:
       - "8020"
       # NOTE: having this port as dynamic one will require modifying
       # `dfs.datanode.address` in `hdfs-site.xml` and probably something
       # else, so using default one for now.
       - "50010:50010"
+      - "50075:50075"
+      - "50070"

tests/func/test_add.py

@@ -288,6 +288,11 @@ def test_add_filtered_files_in_dir(
             "checksum",
             "000002000000000000000000a86fe4d846edc1bf4c355cb6112f141e",
         ),
+        (
+            pytest.lazy_fixture("webhdfs"),
+            "checksum",
+            "000002000000000000000000a86fe4d846edc1bf4c355cb6112f141e00000000",
+        ),
     ],
     indirect=["workspace"],
 )

tests/func/test_api.py

@@ -20,6 +20,7 @@ cloud_names = [
     "http",
     "hdfs",
     "webdav",
+    "webhdfs",
 ]
 clouds = [pytest.lazy_fixture(cloud) for cloud in cloud_names]
 all_clouds = [pytest.lazy_fixture("local_cloud")] + clouds

tests/func/test_data_cloud.py

@@ -30,6 +30,7 @@ all_clouds = [
         "http",
         "hdfs",
         "webdav",
+        "webhdfs",
     ]
 ] + [
     pytest.param(

tests/func/test_gc.py

@@ -356,6 +356,7 @@ def test_gc_not_collect_pipeline_tracked_files(tmp_dir, dvc, run_copy):
         pytest.lazy_fixture("s3"),
         pytest.lazy_fixture("gs"),
         pytest.lazy_fixture("hdfs"),
+        pytest.lazy_fixture("webhdfs"),
         pytest.param(
             pytest.lazy_fixture("ssh"),
             marks=pytest.mark.skipif(

tests/func/test_import_url.py

@@ -122,6 +122,7 @@ def test_import_url_with_no_exec(tmp_dir, dvc, erepo_dir):
         pytest.lazy_fixture("s3"),
         pytest.lazy_fixture("gs"),
         pytest.lazy_fixture("hdfs"),
+        pytest.lazy_fixture("webhdfs"),
         pytest.param(
             pytest.lazy_fixture("ssh"),
             marks=pytest.mark.skipif(

tests/func/test_update.py

@@ -164,6 +164,7 @@ def test_update_before_and_after_dvc_init(tmp_dir, dvc, git_dir):
         pytest.lazy_fixture("s3"),
         pytest.lazy_fixture("gs"),
         pytest.lazy_fixture("hdfs"),
+        pytest.lazy_fixture("webhdfs"),
         pytest.param(
             pytest.lazy_fixture("ssh"),
             marks=pytest.mark.skipif(

tests/remotes/__init__.py

@@ -3,7 +3,7 @@ import subprocess
 import pytest
 
 from .azure import Azure, azure, azure_server  # noqa: F401
-from .hdfs import HDFS, hadoop, hdfs, hdfs_server  # noqa: F401
+from .hdfs import HDFS, hadoop, hdfs, hdfs_server, webhdfs  # noqa: F401
 from .http import HTTP, http, http_server  # noqa: F401
 from .local import Local, local_cloud, local_remote  # noqa: F401
 from .oss import (  # noqa: F401

tests/remotes/hdfs.py

@@ -119,6 +119,7 @@ def hdfs_server(hadoop, docker_compose, docker_services):
     import pyarrow
 
     port = docker_services.port_for("hdfs", 8020)
+    web_port = docker_services.port_for("hdfs", 50070)
 
     def _check():
         try:
@@ -136,11 +137,70 @@ def hdfs_server(hadoop, docker_compose, docker_services):
 
     docker_services.wait_until_responsive(timeout=30.0, pause=5, check=_check)
 
-    return port
+    return {"hdfs": port, "webhdfs": web_port}
 
 
 @pytest.fixture
 def hdfs(hdfs_server):
-    port = hdfs_server
+    port = hdfs_server["hdfs"]
     url = f"hdfs://127.0.0.1:{port}/{uuid.uuid4()}"
     yield HDFS(url)
+
+
+class WebHDFS(Base, URLInfo):  # pylint: disable=abstract-method
+    @contextmanager
+    def _webhdfs(self):
+        from hdfs import InsecureClient
+
+        client = InsecureClient(f"http://{self.host}:{self.port}", self.user)
+        yield client
+
+    def is_file(self):
+        with self._webhdfs() as _hdfs:
+            return _hdfs.status(self.path)["type"] == "FILE"
+
+    def is_dir(self):
+        with self._webhdfs() as _hdfs:
+            return _hdfs.status(self.path)["type"] == "DIRECTORY"
+
+    def exists(self):
+        with self._webhdfs() as _hdfs:
+            return _hdfs.status(self.path, strict=False) is not None
+
+    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
+        assert mode == 0o777
+        assert parents
+        assert not exist_ok
+
+        with self._webhdfs() as _hdfs:
+            # NOTE: hdfs.makekdirs always creates parents
+            _hdfs.makedirs(self.path, permission=mode)
+
+    def write_bytes(self, contents):
+        with self._webhdfs() as _hdfs:
+            with _hdfs.write(self.path, overwrite=True) as writer:
+                writer.write(contents)
+
+    def write_text(self, contents, encoding=None, errors=None):
+        if not encoding:
+            encoding = locale.getpreferredencoding(False)
+        assert errors is None
+        self.write_bytes(contents.encode(encoding))
+
+    def read_bytes(self):
+        with self._webhdfs() as _hdfs:
+            with _hdfs.read(self.path) as reader:
+                return reader.read()
+
+    def read_text(self, encoding=None, errors=None):
+        if not encoding:
+            encoding = locale.getpreferredencoding(False)
+        assert errors is None
+        return self.read_bytes().decode(encoding)
+
+
+@pytest.fixture
+def webhdfs(hdfs_server):
+    port = hdfs_server["webhdfs"]
+    url = f"webhdfs://127.0.0.1:{port}/{uuid.uuid4()}"
+    yield WebHDFS(url)

tests/unit/remote/test_webhdfs.py

@@ -0,0 +1,24 @@
+from dvc.tree.webhdfs import WebHDFSTree
+
+user = "test"
+webhdfs_token = "token"
+webhdfs_alias = "alias-name"
+hdfscli_config = "path/to/cli/config"
+
+
+def test_init(dvc):
+    url = "webhdfs://[email protected]:50070"
+    config = {
+        "url": url,
+        "webhdfs_token": webhdfs_token,
+        "webhdfs_alias": webhdfs_alias,
+        "hdfscli_config": hdfscli_config,
+        "user": user,
+    }
+
+    tree = WebHDFSTree(dvc, config)
+    assert tree.path_info == url
+    assert tree.token == webhdfs_token
+    assert tree.alias == webhdfs_alias
+    assert tree.path_info.user == user
+    assert tree.hdfscli_config == hdfscli_config