From d038bca05eeb78fdf0b9718db7b9a3789790900e Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Tue, 19 Jan 2021 20:03:16 +0200 Subject: [PATCH] tests: add basic [web]hdfs mockers (#5279) --- dvc/tree/hdfs.py | 86 +++++++------- dvc/tree/webhdfs.py | 9 +- setup.py | 1 + tests/remotes/__init__.py | 3 +- tests/remotes/hdfs.py | 141 +++++++++++++++------- tests/remotes/webhdfs.py | 223 +++++++++++++++++++++++++++++++++++ tests/unit/tree/test_hdfs.py | 30 +++++ 7 files changed, 400 insertions(+), 93 deletions(-) create mode 100644 tests/remotes/webhdfs.py create mode 100644 tests/unit/tree/test_hdfs.py diff --git a/dvc/tree/hdfs.py b/dvc/tree/hdfs.py index 1250ebfe6..5f50add7a 100644 --- a/dvc/tree/hdfs.py +++ b/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.*)" + 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.*)" - 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 diff --git a/dvc/tree/webhdfs.py b/dvc/tree/webhdfs.py index 816748ce4..9e3454470 100644 --- a/dvc/tree/webhdfs.py +++ b/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: diff --git a/setup.py b/setup.py index d8343b4be..d795a8e3d 100644 --- a/setup.py +++ b/setup.py @@ -151,6 +151,7 @@ tests_requirements = [ "filelock", "mypy", "wsgidav", + "crc32c", ] setup( diff --git a/tests/remotes/__init__.py b/tests/remotes/__init__.py index d075923d3..320665d94 100644 --- a/tests/remotes/__init__.py +++ b/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 = { diff --git a/tests/remotes/hdfs.py b/tests/remotes/hdfs.py index 35963a205..07193373a 100644 --- a/tests/remotes/hdfs.py +++ b/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) diff --git a/tests/remotes/webhdfs.py b/tests/remotes/webhdfs.py new file mode 100644 index 000000000..0e5d97e9f --- /dev/null +++ b/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) diff --git a/tests/unit/tree/test_hdfs.py b/tests/unit/tree/test_hdfs.py new file mode 100644 index 000000000..aaa901ba9 --- /dev/null +++ b/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