Add setOwner support (#117)

Author: Kimahriman

python/python/hdfs_native/__init__.py

@@ -1,5 +1,5 @@
 import io
-from typing import Iterator
+from typing import Iterator, Optional
 from typing_extensions import Buffer
 
 from ._internal import *
@@ -113,3 +113,9 @@ def set_times(self, path: str, mtime: int, atime: int) -> None:
         Changes the modification time and access time of the file at `path` to `mtime` and `atime`, respectively.
         """
         return self.inner.set_times(path, mtime, atime)
+
+    def set_owner(self, path: str, owner: Optional[str] = None, group: Optional[str] = None) -> None:
+        """
+        Sets the owner and/or group for the file at `path`
+        """
+        return self.inner.set_owner(path, owner, group)

python/python/hdfs_native/_internal.pyi

@@ -55,4 +55,6 @@ class RawClient:
 
     def delete(self, path: str, recursive: bool) -> bool: ...
 
-    def set_times(self, path: str, mtime: int, atime: int) -> None: ...
+    def set_times(self, path: str, mtime: int, atime: int) -> None: ...
+
+    def set_owner(self, path: str, owner: Optional[str], group: Optional[str]) -> None: ...

python/src/lib.rs

@@ -243,6 +243,15 @@ impl RawClient {
     pub fn set_times(&self, path: &str, mtime: u64, atime: u64) -> PyHdfsResult<()> {
         Ok(self.rt.block_on(self.inner.set_times(path, mtime, atime))?)
     }
+
+    pub fn set_owner(
+        &self,
+        path: &str,
+        owner: Option<&str>,
+        group: Option<&str>,
+    ) -> PyHdfsResult<()> {
+        Ok(self.rt.block_on(self.inner.set_owner(path, owner, group))?)
+    }
 }
 
 /// A Python module implemented in Rust.

python/tests/test_integration.py

@@ -59,5 +59,20 @@ def test_integration(minidfs: str):
     assert file_info.modification_time == mtime
     assert file_info.access_time == atime
 
-    client.delete("/testfile", False)
 
+    client.set_owner("/testfile", "testuser", "testgroup")
+    file_info = client.get_file_info("/testfile")
+    assert file_info.owner == "testuser"
+    assert file_info.group == "testgroup"
+
+    client.set_owner("/testfile", owner="testuser2")
+    file_info = client.get_file_info("/testfile")
+    assert file_info.owner == "testuser2"
+    assert file_info.group == "testgroup"
+
+    client.set_owner("/testfile", group="testgroup2")
+    file_info = client.get_file_info("/testfile")
+    assert file_info.owner == "testuser2"
+    assert file_info.group == "testgroup2"
+
+    client.delete("/testfile", False)

rust/src/client.rs

@@ -431,6 +431,20 @@ impl Client {
             .await?;
         Ok(())
     }
+
+    /// Optionally sets the owner and group for a file.
+    pub async fn set_owner(
+        &self,
+        path: &str,
+        owner: Option<&str>,
+        group: Option<&str>,
+    ) -> Result<()> {
+        let (link, resolved_path) = self.mount_table.resolve(path);
+        link.protocol
+            .set_owner(&resolved_path, owner, group)
+            .await?;
+        Ok(())
+    }
 }
 
 impl Default for Client {

rust/src/hdfs/protocol.rs

@@ -411,6 +411,30 @@ impl NamenodeProtocol {
         debug!("setTimes response: {:?}", &decoded);
         Ok(decoded)
     }
+
+    pub(crate) async fn set_owner(
+        &self,
+        src: &str,
+        owner: Option<&str>,
+        group: Option<&str>,
+    ) -> Result<hdfs::SetOwnerResponseProto> {
+        let message = hdfs::SetOwnerRequestProto {
+            src: src.to_string(),
+            username: owner.map(str::to_string),
+            groupname: group.map(str::to_string),
+        };
+
+        debug!("setOwner request: {:?}", &message);
+
+        let response = self
+            .proxy
+            .call("setOwner", message.encode_length_delimited_to_vec())
+            .await?;
+
+        let decoded = hdfs::SetOwnerResponseProto::decode_length_delimited(response)?;
+        debug!("setOwner response: {:?}", &decoded);
+        Ok(decoded)
+    }
 }
 
 impl Drop for NamenodeProtocol {

rust/tests/test_integration.rs

@@ -188,6 +188,7 @@ mod test {
         // We use writing to create files, so do this after
         test_recursive_listing(&client).await?;
         test_set_times(&client).await?;
+        test_set_owner(&client).await?;
 
         Ok(())
     }
@@ -348,6 +349,38 @@ mod test {
         assert_eq!(file_info.modification_time, mtime);
         assert_eq!(file_info.access_time, atime);
 
+        client.delete("/test", false).await?;
+
+        Ok(())
+    }
+
+    async fn test_set_owner(client: &Client) -> Result<()> {
+        client
+            .create("/test", WriteOptions::default())
+            .await?
+            .close()
+            .await?;
+
+        client
+            .set_owner("/test", Some("testuser"), Some("testgroup"))
+            .await?;
+        let file_info = client.get_file_info("/test").await?;
+
+        assert_eq!(file_info.owner, "testuser");
+        assert_eq!(file_info.group, "testgroup");
+
+        client.set_owner("/test", Some("testuser2"), None).await?;
+        let file_info = client.get_file_info("/test").await?;
+
+        assert_eq!(file_info.owner, "testuser2");
+        assert_eq!(file_info.group, "testgroup");
+
+        client.set_owner("/test", None, Some("testgroup2")).await?;
+        let file_info = client.get_file_info("/test").await?;
+
+        assert_eq!(file_info.owner, "testuser2");
+        assert_eq!(file_info.group, "testgroup2");
+
         Ok(())
     }
 }