diff --git a/.pylintrc b/.pylintrc index e036b6217..85769da0d 100644 --- a/.pylintrc +++ b/.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). diff --git a/dvc/cache/__init__.py b/dvc/cache/__init__.py index fc78f47cd..dd6471983 100644 --- a/dvc/cache/__init__.py +++ b/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 diff --git a/dvc/config.py b/dvc/config.py index d3f31e9d8..bb4e31a7d 100644 --- a/dvc/config.py +++ b/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, diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index 7ed45fd96..22db00bae 100644 --- a/dvc/dependency/__init__.py +++ b/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, } diff --git a/dvc/dependency/webhdfs.py b/dvc/dependency/webhdfs.py new file mode 100644 index 000000000..e3d5f1deb --- /dev/null +++ b/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 diff --git a/dvc/output/__init__.py b/dvc/output/__init__.py index 0ab8920f2..9b230015d 100644 --- a/dvc/output/__init__.py +++ b/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() diff --git a/dvc/output/webhdfs.py b/dvc/output/webhdfs.py new file mode 100644 index 000000000..7fd28229e --- /dev/null +++ b/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 diff --git a/dvc/scheme.py b/dvc/scheme.py index 76c6d7a49..3c177dcd3 100644 --- a/dvc/scheme.py +++ b/dvc/scheme.py @@ -1,6 +1,7 @@ class Schemes: SSH = "ssh" HDFS = "hdfs" + WEBHDFS = "webhdfs" S3 = "s3" AZURE = "azure" HTTP = "http" diff --git a/dvc/tree/__init__.py b/dvc/tree/__init__.py index cbd1595a6..eacb12352 100644 --- a/dvc/tree/__init__.py +++ b/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 ] diff --git a/dvc/tree/webhdfs.py b/dvc/tree/webhdfs.py new file mode 100644 index 000000000..f6e77ef90 --- /dev/null +++ b/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) + ) diff --git a/setup.py b/setup.py index 0df9c255c..66d413836 100644 --- a/setup.py +++ b/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, }, diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index 04d677485..3bf4e45cd 100644 --- a/tests/docker-compose.yml +++ b/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" diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 483cdb857..b2dcd38c8 100644 --- a/tests/func/test_add.py +++ b/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"], ) diff --git a/tests/func/test_api.py b/tests/func/test_api.py index 92955007f..4e3301ef2 100644 --- a/tests/func/test_api.py +++ b/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 diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 9ef36eb0c..888546030 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -30,6 +30,7 @@ all_clouds = [ "http", "hdfs", "webdav", + "webhdfs", ] ] + [ pytest.param( diff --git a/tests/func/test_gc.py b/tests/func/test_gc.py index 90b5664a4..4422f207f 100644 --- a/tests/func/test_gc.py +++ b/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( diff --git a/tests/func/test_import_url.py b/tests/func/test_import_url.py index f1eb4b83f..906566bea 100644 --- a/tests/func/test_import_url.py +++ b/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( diff --git a/tests/func/test_update.py b/tests/func/test_update.py index 4dd205f29..26fce3e9d 100644 --- a/tests/func/test_update.py +++ b/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( diff --git a/tests/remotes/__init__.py b/tests/remotes/__init__.py index 1750dcef1..5e4d8b9d5 100644 --- a/tests/remotes/__init__.py +++ b/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 diff --git a/tests/remotes/hdfs.py b/tests/remotes/hdfs.py index 7a765f080..45d687f20 100644 --- a/tests/remotes/hdfs.py +++ b/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) diff --git a/tests/unit/remote/test_webhdfs.py b/tests/unit/remote/test_webhdfs.py new file mode 100644 index 000000000..e25431225 --- /dev/null +++ b/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://test@127.0.0.1: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