From c5b1a73dc5f6cabf19c4e1deb31feed0dc9de223 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gu=C3=B0mundur=20Heimisson?= <gudmundur.heimisson@gmail.com>
Date: Sun, 7 Nov 2021 19:42:10 +0100
Subject: [PATCH] webhdfs: expose kerberos and https options (#6936)

* Add fsspec options for webhdfs

* Add requests-kerberos dependency

* Update dvc/fs/webhdfs.py

Simplify config.pop.

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>

* Update dvc/fs/webhdfs.py

* Simplify kerb_kwargs config setting.

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/config_schema.py              | 10 +++---
 dvc/fs/webhdfs.py                 | 11 +++---
 setup.cfg                         |  1 +
 tests/unit/remote/test_webhdfs.py | 59 ++++++++++++++++++++++++-------
 4 files changed, 60 insertions(+), 21 deletions(-)

diff --git a/dvc/config_schema.py b/dvc/config_schema.py
index 42c5ee37f..7427825ca 100644
--- a/dvc/config_schema.py
+++ b/dvc/config_schema.py
@@ -179,10 +179,12 @@ SCHEMA = {
                 },
                 "hdfs": {"user": str, "kerb_ticket": str, **REMOTE_COMMON},
                 "webhdfs": {
-                    "hdfscli_config": str,
-                    "webhdfs_token": str,
-                    "user": str,
-                    "webhdfs_alias": str,
+                    "kerberos": Bool,
+                    "kerberos_principal": str,
+                    "proxy_to": str,
+                    "ssl_verify": Any(Bool, str),
+                    "token": str,
+                    "use_https": Bool,
                     **REMOTE_COMMON,
                 },
                 "azure": {
diff --git a/dvc/fs/webhdfs.py b/dvc/fs/webhdfs.py
index 964b5b3fb..9a51c9cbd 100644
--- a/dvc/fs/webhdfs.py
+++ b/dvc/fs/webhdfs.py
@@ -31,9 +31,10 @@ class WebHDFSFileSystem(CallbackMixin, FSSpecWrapper):
         )
 
     def _prepare_credentials(self, **config):
-        if "webhdfs_token" in config:
-            config["token"] = config.pop("webhdfs_token")
-
+        self._ssl_verify = config.pop("ssl_verify", True)
+        principal = config.pop("kerberos_principal", None)
+        if principal:
+            config["kerb_kwargs"] = {"principal": principal}
         return config
 
     @wrap_prop(threading.Lock())
@@ -41,7 +42,9 @@ class WebHDFSFileSystem(CallbackMixin, FSSpecWrapper):
     def fs(self):
         from fsspec.implementations.webhdfs import WebHDFS
 
-        return WebHDFS(**self.fs_args)
+        fs = WebHDFS(**self.fs_args)
+        fs.session.verify = self._ssl_verify
+        return fs
 
     def checksum(self, path_info):
         path = self._with_bucket(path_info)
diff --git a/setup.cfg b/setup.cfg
index 114a05d82..f492b259b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -112,6 +112,7 @@ ssh_gssapi = sshfs[gssapi]>=2021.8.1
 webdav = webdav4>=0.9.3
 # not to break `dvc[webhdfs]`
 webhdfs =
+    requests-kerberos==0.13.0
 terraform = tpi[ssh]>=2.1.0
 tests =
     %(terraform)s
diff --git a/tests/unit/remote/test_webhdfs.py b/tests/unit/remote/test_webhdfs.py
index b267c32fa..341aa4928 100644
--- a/tests/unit/remote/test_webhdfs.py
+++ b/tests/unit/remote/test_webhdfs.py
@@ -1,21 +1,54 @@
+from unittest.mock import Mock, create_autospec
+
+import pytest
+import requests
+
 from dvc.fs.webhdfs import WebHDFSFileSystem
 
+host = "host"
+kerberos = False
+kerberos_principal = "principal"
+port = 12345
+proxy_to = "proxy"
+ssl_verify = False
+token = "token"
+use_https = True
 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 = {
-        "host": url,
-        "webhdfs_token": webhdfs_token,
-        "webhdfs_alias": webhdfs_alias,
-        "hdfscli_config": hdfscli_config,
-        "user": user,
+@pytest.fixture()
+def webhdfs_config():
+    url = f"webhdfs://{user}@{host}:{port}"
+    url_config = WebHDFSFileSystem._get_kwargs_from_urls(url)
+    return {
+        "kerberos": kerberos,
+        "kerberos_principal": kerberos_principal,
+        "proxy_to": proxy_to,
+        "ssl_verify": ssl_verify,
+        "token": token,
+        "use_https": use_https,
+        **url_config,
     }
 
-    fs = WebHDFSFileSystem(**config)
-    assert fs.fs_args["token"] == webhdfs_token
+
+def test_init(dvc, webhdfs_config):
+    fs = WebHDFSFileSystem(**webhdfs_config)
+    assert fs.fs_args["host"] == host
+    assert fs.fs_args["token"] == token
     assert fs.fs_args["user"] == user
+    assert fs.fs_args["port"] == port
+    assert fs.fs_args["kerberos"] == kerberos
+    assert fs.fs_args["kerb_kwargs"] == {"principal": kerberos_principal}
+    assert fs.fs_args["proxy_to"] == proxy_to
+    assert fs.fs_args["use_https"] == use_https
+
+
+def test_verify_ssl(dvc, webhdfs_config, monkeypatch):
+    mock_session = create_autospec(requests.Session)
+    monkeypatch.setattr(requests, "Session", Mock(return_value=mock_session))
+    # can't have token at the same time as user or proxy_to
+    del webhdfs_config["token"]
+    fs = WebHDFSFileSystem(**webhdfs_config)
+    # ssl verify can't be set until after the file system is instantiated
+    fs.fs  # pylint: disable=pointless-statement
+    assert mock_session.verify == ssl_verify