dvc/tests/remotes/webhdfs.py

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

        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)