tests: add basic [web]hdfs mockers (#5279)

dvc/tree/hdfs.py

@@ -19,6 +19,45 @@ from .pool import get_connection
 logger = logging.getLogger(__name__)
 
 
+def _hadoop_fs(cmd, user=None):
+    cmd = "hadoop fs -" + cmd
+    if user:
+        cmd = f"HADOOP_USER_NAME={user} " + cmd
+
+    # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
+    # See https://github.com/iterative/dvc/issues/1197.
+    close_fds = os.name != "nt"
+
+    executable = os.getenv("SHELL") if os.name != "nt" else None
+    p = subprocess.Popen(
+        cmd,
+        shell=True,
+        close_fds=close_fds,
+        executable=executable,
+        env=fix_env(os.environ),
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    out, err = p.communicate()
+    if p.returncode != 0:
+        raise RemoteCmdError("hdfs", cmd, p.returncode, err)
+    return out.decode("utf-8")
+
+
+def _group(regex, s, gname):
+    match = re.match(regex, s)
+    assert match is not None
+    return match.group(gname)
+
+
+def _hadoop_fs_checksum(path_info):
+    # NOTE: pyarrow doesn't support checksum, so we need to use hadoop
+    regex = r".*\t.*\t(?P<checksum>.*)"
+    stdout = _hadoop_fs(f"checksum {path_info.url}", user=path_info.user)
+    return _group(regex, stdout, "checksum")
+
+
 class HDFSTree(BaseTree):
     scheme = Schemes.HDFS
     REQUIRES = {"pyarrow": "pyarrow"}
@@ -188,52 +227,13 @@ class HDFSTree(BaseTree):
             file_info = hdfs.get_file_info(path_info.path)
             return file_info.size
 
-    def hadoop_fs(self, cmd, user=None):
-        cmd = "hadoop fs -" + cmd
-        if user:
-            cmd = f"HADOOP_USER_NAME={user} " + cmd
-
-        # NOTE: close_fds doesn't work with redirected stdin/stdout/stderr.
-        # See https://github.com/iterative/dvc/issues/1197.
-        close_fds = os.name != "nt"
-
-        executable = os.getenv("SHELL") if os.name != "nt" else None
-        p = subprocess.Popen(
-            cmd,
-            shell=True,
-            close_fds=close_fds,
-            executable=executable,
-            env=fix_env(os.environ),
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-        out, err = p.communicate()
-        if p.returncode != 0:
-            raise RemoteCmdError(self.scheme, cmd, p.returncode, err)
-        return out.decode("utf-8")
-
-    @staticmethod
-    def _group(regex, s, gname):
-        match = re.match(regex, s)
-        assert match is not None
-        return match.group(gname)
-
     def get_file_hash(self, path_info):
         # NOTE: pyarrow doesn't support checksum, so we need to use hadoop
-        regex = r".*\t.*\t(?P<checksum>.*)"
-        stdout = self.hadoop_fs(
-            f"checksum {path_info.url}", user=path_info.user
+        return HashInfo(
+            self.PARAM_CHECKSUM,
+            _hadoop_fs_checksum(path_info),
+            size=self.getsize(path_info),
         )
-        hash_info = HashInfo(
-            self.PARAM_CHECKSUM, self._group(regex, stdout, "checksum"),
-        )
-
-        with self.hdfs(path_info) as hdfs:
-            file_info = hdfs.get_file_info(path_info.path)
-            hash_info.size = file_info.size
-
-        return hash_info
 
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs

dvc/tree/webhdfs.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import posixpath
 import shutil
 import threading
 from contextlib import contextmanager
@@ -34,6 +35,7 @@ class WebHDFSTree(BaseTree):
     PATH_CLS = CloudURLInfo
     REQUIRES = {"hdfs": "hdfs"}
     PARAM_CHECKSUM = "checksum"
+    TRAVERSE_PREFIX_LEN = 2
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
@@ -105,10 +107,9 @@ class WebHDFSTree(BaseTree):
             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)
+        for path, _, fnames in self.hdfs_client.walk(root):
+            for fname in fnames:
+                yield path_info.replace(path=posixpath.join(path, fname))
 
     def remove(self, path_info):
         if path_info.scheme != self.scheme:

setup.py

@@ -151,6 +151,7 @@ tests_requirements = [
     "filelock",
     "mypy",
     "wsgidav",
+    "crc32c",
 ]
 
 setup(

tests/remotes/__init__.py

@@ -14,7 +14,7 @@ from .gs import (  # noqa: F401; noqa: F401
     TEST_GCP_REPO_BUCKET,
     gs,
 )
-from .hdfs import HDFS, hadoop, hdfs, hdfs_server, webhdfs  # noqa: F401
+from .hdfs import HDFS, hadoop, hdfs, hdfs_server, real_hdfs  # 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
@@ -32,6 +32,7 @@ from .ssh import (  # noqa: F401; noqa: F401
     ssh_server,
 )
 from .webdav import Webdav, webdav, webdav_server  # noqa: F401
+from .webhdfs import WebHDFS, real_webhdfs, webhdfs  # noqa: F401
 
 TEST_REMOTE = "upstream"
 TEST_CONFIG = {

tests/remotes/hdfs.py

@@ -4,6 +4,8 @@ import platform
 import subprocess
 import uuid
 from contextlib import contextmanager
+from pathlib import Path
+from tempfile import TemporaryDirectory
 
 import pytest
 
@@ -12,6 +14,8 @@ from tests import PY39, PYARROW_NOT_AVAILABLE
 
 from .base import Base
 
+_hdfs_root = TemporaryDirectory()
+
 
 class HDFS(Base, URLInfo):  # pylint: disable=abstract-method
     @contextmanager
@@ -154,66 +158,113 @@ def hdfs_server(hadoop, docker_compose, docker_services):
 
 
 @pytest.fixture
-def hdfs(hdfs_server):
+def real_hdfs(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
+def md5md5crc32c(path):
+    # https://github.com/colinmarc/hdfs/blob/f2f512db170db82ad41590c4ba3b7718b13317d2/file_reader.go#L76
+    import hashlib
 
-        client = InsecureClient(f"http://{self.host}:{self.port}", self.user)
-        yield client
+    from crc32c import crc32c  # pylint: disable=no-name-in-module
 
-    def is_file(self):
-        with self._webhdfs() as _hdfs:
-            return _hdfs.status(self.path)["type"] == "FILE"
+    # dfs.bytes-per-checksum = 512, default on hadoop 2.7
+    bytes_per_checksum = 512
+    padded = 32
+    total = 0
 
-    def is_dir(self):
-        with self._webhdfs() as _hdfs:
-            return _hdfs.status(self.path)["type"] == "DIRECTORY"
+    md5md5 = hashlib.md5()
 
-    def exists(self):
-        with self._webhdfs() as _hdfs:
-            return _hdfs.status(self.path, strict=False) is not None
+    with open(path, "rb") as fobj:
+        while True:
+            block = fobj.read(bytes_per_checksum)
+            if not block:
+                break
 
-    def mkdir(self, mode=0o777, parents=False, exist_ok=False):
-        assert mode == 0o777
-        assert parents
-        assert not exist_ok
+            crc_int = crc32c(block)
 
-        with self._webhdfs() as _hdfs:
-            # NOTE: hdfs.makekdirs always creates parents
-            _hdfs.makedirs(self.path, permission=mode)
+            # NOTE: hdfs is big-endian
+            crc_bytes = crc_int.to_bytes(
+                (crc_int.bit_length() + 7) // 8, "big"
+            )
 
-    def write_bytes(self, contents):
-        with self._webhdfs() as _hdfs:
-            with _hdfs.write(self.path, overwrite=True) as writer:
-                writer.write(contents)
+            md5 = hashlib.md5(crc_bytes).digest()
 
-    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))
+            total += len(md5)
+            if padded < total:
+                padded *= 2
 
-    def read_bytes(self):
-        with self._webhdfs() as _hdfs:
-            with _hdfs.read(self.path) as reader:
-                return reader.read()
+            md5md5.update(md5)
 
-    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)
+    md5md5.update(b"\0" * (padded - total))
+    return "000002000000000000000000" + md5md5.hexdigest()
+
+
+def hadoop_fs_checksum(path_info):
+
+    return md5md5crc32c(Path(_hdfs_root.name) / path_info.path.lstrip("/"))
+
+
+class FakeHadoopFileSystem:
+    def __init__(self, *args, **kwargs):
+        from pyarrow.fs import LocalFileSystem
+
+        self._root = Path(_hdfs_root.name)
+        self._fs = LocalFileSystem()
+
+    def _path(self, path):
+        from pyarrow.fs import FileSelector
+
+        if isinstance(path, FileSelector):
+            return FileSelector(
+                os.fspath(self._root / path.base_dir.lstrip("/")),
+                path.allow_not_found,
+                path.recursive,
+            )
+
+        return os.fspath(self._root / path.lstrip("/"))
+
+    def create_dir(self, path):
+        return self._fs.create_dir(self._path(path))
+
+    def open_input_stream(self, path):
+        return self._fs.open_input_stream(self._path(path))
+
+    def open_output_stream(self, path):
+        import posixpath
+
+        # NOTE: HadoopFileSystem.open_output_stream creates directories
+        # automatically.
+        self.create_dir(posixpath.dirname(path))
+        return self._fs.open_output_stream(self._path(path))
+
+    def get_file_info(self, path):
+        return self._fs.get_file_info(self._path(path))
+
+    def move(self, from_path, to_path):
+        self._fs.move(self._path(from_path), self._path(to_path))
+
+    def delete_file(self, path):
+        self._fs.delete_file(self._path(path))
 
 
 @pytest.fixture
-def webhdfs(hdfs_server):
-    port = hdfs_server["webhdfs"]
-    url = f"webhdfs://127.0.0.1:{port}/{uuid.uuid4()}"
-    yield WebHDFS(url)
+def hdfs(mocker):
+    if PY39:
+        pytest.skip(PYARROW_NOT_AVAILABLE)
+
+    # Windows might not have Visual C++ Redistributable for Visual Studio
+    # 2015 installed, which will result in the following error:
+    # "The pyarrow installation is not built with support for
+    # 'HadoopFileSystem'"
+    mocker.patch("pyarrow.fs._not_imported", [])
+    mocker.patch(
+        "pyarrow.fs.HadoopFileSystem", FakeHadoopFileSystem, create=True
+    )
+
+    mocker.patch("dvc.tree.hdfs._hadoop_fs_checksum", hadoop_fs_checksum)
+
+    url = f"hdfs://example.com:12345/{uuid.uuid4()}"
+    yield HDFS(url)

tests/remotes/webhdfs.py

@@ -0,0 +1,223 @@
+import locale
+import os
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+
+import pytest
+
+from dvc.path_info import URLInfo
+
+from .base import Base
+from .hdfs import _hdfs_root, md5md5crc32c
+
+
+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 real_webhdfs(hdfs_server):
+    port = hdfs_server["webhdfs"]
+    url = f"webhdfs://127.0.0.1:{port}/{uuid.uuid4()}"
+    yield WebHDFS(url)
+
+
+class FakeClient:
+    def __init__(self, *args, **kwargs):
+        self._root = Path(_hdfs_root.name)
+
+    def _path(self, path):
+        return self._root / path.lstrip("/")
+
+    def makedirs(self, path, permission=None):
+        self._path(path).mkdir(
+            mode=permission or 0o777, exist_ok=True, parents=True
+        )
+
+    def write(self, hdfs_path, overwrite=False):
+        from hdfs.util import HdfsError
+
+        path = self._path(hdfs_path)
+
+        if not overwrite and path.exists():
+            raise HdfsError(f"Remote path {hdfs_path} already exists.")
+
+        path.parent.mkdir(parents=True, exist_ok=True)
+        return path.open("wb")
+
+    @contextmanager
+    def read(
+        self,
+        hdfs_path,
+        encoding=None,
+        chunk_size=0,
+        delimiter=None,
+        progress=None,
+    ):
+        pobj = self._path(hdfs_path)
+        if not chunk_size and not delimiter:
+            if encoding:
+                yield pobj.open("r", encoding=encoding)
+            else:
+                yield pobj.open("rb")
+        else:
+            if delimiter:
+                data = pobj.open("r", encoding=encoding, newline=delimiter)
+            else:
+
+                def read_chunks(fobj, _chunk_size):
+                    while True:
+                        chunk = fobj.read(_chunk_size)
+                        if not chunk:
+                            break
+                        yield chunk
+
+                data = read_chunks(
+                    pobj.open("rb", encoding=encoding), chunk_size
+                )
+
+            if progress:
+
+                def reader(_hdfs_path, _progress):
+                    nbytes = 0
+                    for chunk in data:
+                        nbytes += len(chunk)
+                        _progress(_hdfs_path, nbytes)
+                        yield chunk
+                    _progress(_hdfs_path, -1)
+
+                yield reader(hdfs_path, progress)
+            else:
+                yield data
+
+    def walk(self, hdfs_path):
+        import posixpath
+
+        local_path = self._path(hdfs_path)
+        for local_root, dnames, fnames in os.walk(local_path):
+            if local_root == os.fspath(local_path):
+                root = hdfs_path
+            else:
+                root = posixpath.join(
+                    hdfs_path, os.path.relpath(local_root, local_path)
+                )
+            yield (
+                root,
+                dnames,
+                fnames,
+            )
+
+    def delete(self, hdfs_path):
+        return self._path(hdfs_path).unlink()
+
+    def status(self, hdfs_path, strict=True):
+        from hdfs.util import HdfsError
+
+        try:
+            return {"length": self._path(hdfs_path).stat().st_size}
+        except FileNotFoundError:
+            if not strict:
+                return None
+            raise HdfsError(
+                f"File does not exist: {hdfs_path}",
+                exception="FileNotFoundException",
+            )
+
+    def checksum(self, hdfs_path):
+        return {
+            "algorithm": "MD5-of-0MD5-of-512CRC32",
+            "bytes": md5md5crc32c(self._path(hdfs_path)) + "00000000",
+            "size": 28,
+        }
+
+    def rename(self, from_path, to_path):
+        from dvc.utils.fs import move
+
+        move(self._path(from_path), self._path(to_path))
+
+    def upload(
+        self,
+        hdfs_path,
+        local_path,
+        chunk_size=2 ** 16,
+        progress=None,
+        **kwargs,
+    ):
+        with open(local_path, "rb") as from_fobj:
+            with self.write(hdfs_path, **kwargs) as to_fobj:
+                nbytes = 0
+                while True:
+                    chunk = from_fobj.read(chunk_size)
+                    if not chunk:
+                        break
+                    if progress:
+                        nbytes += len(chunk)
+                        progress(local_path, nbytes)
+                    to_fobj.write(chunk)
+
+    def download(self, hdfs_path, local_path, **kwargs):
+        from dvc.utils.fs import makedirs
+
+        kwargs.setdefault("chunk_size", 2 ** 16)
+
+        makedirs(os.path.dirname(local_path), exist_ok=True)
+        with open(local_path, "wb") as writer:
+            with self.read(hdfs_path, **kwargs) as reader:
+                for chunk in reader:
+                    writer.write(chunk)
+
+
+@pytest.fixture
+def webhdfs(mocker):
+    mocker.patch("hdfs.InsecureClient", FakeClient)
+    url = f"webhdfs://example.com:12345/{uuid.uuid4()}"
+    yield WebHDFS(url)

tests/unit/tree/test_hdfs.py

@@ -0,0 +1,30 @@
+import os
+
+from dvc.path_info import URLInfo
+from dvc.tree.hdfs import _hadoop_fs_checksum
+
+
+def test_hadoop_fs_checksum(mocker):
+    mock_proc = mocker.Mock()
+
+    out = b"/path/to/file\tMD5-of-0MD5-of-512CRC32C\t123456789"
+    err = b""
+    mock_proc.configure_mock(
+        **{"communicate.return_value": (out, err), "returncode": 0}
+    )
+    mock_popen = mocker.patch("subprocess.Popen", return_value=mock_proc)
+
+    path_info = URLInfo("hdfs://example.com:1234/path/to/file")
+
+    assert _hadoop_fs_checksum(path_info) == "123456789"
+    mock_popen.assert_called_once_with(
+        "hadoop fs -checksum hdfs://example.com:1234/path/to/file",
+        shell=True,
+        close_fds=os.name != "nt",
+        executable=os.getenv("SHELL") if os.name != "nt" else None,
+        env=os.environ,
+        stdin=-1,
+        stdout=-1,
+        stderr=-1,
+    )
+    assert mock_proc.communicate.called