Add delete_object support in s3torchconnectorclient (#129)

- Add delete_object support in s3torchconnectorclient

---------

Co-authored-by: Simon Beal <[email protected]>

Author: muddyfish

s3torchconnectorclient/python/src/s3torchconnectorclient/_mountpoint_s3_client.pyi

@@ -30,6 +30,7 @@ class MountpointS3Client:
         self, bucket: str, prefix: str = "", delimiter: str = "", max_keys: int = 1000
     ) -> ListObjectStream: ...
     def head_object(self, bucket: str, key: str) -> ObjectInfo: ...
+    def delete_object(self, bucket: str, key: str) -> None: ...
 
 class MockMountpointS3Client:
     def __init__(

s3torchconnectorclient/python/tst/integration/conftest.py

@@ -107,3 +107,8 @@ def put_object_tests_directory(request) -> BucketPrefixFixture:
 @pytest.fixture
 def checkpoint_directory(request) -> BucketPrefixFixture:
     return get_test_bucket_prefix(f"{request.node.name}/checkpoint_directory")
+
+
+@pytest.fixture
+def empty_directory(request) -> BucketPrefixFixture:
+    return get_test_bucket_prefix(f"{request.node.name}/empty_directory")

s3torchconnectorclient/python/tst/integration/test_mountpoint_s3_integration.py

@@ -258,6 +258,37 @@ def test_head_object(sample_directory):
         assert object_info.etag == expected_etag
 
 
+def test_delete_object(sample_directory):
+    client = MountpointS3Client(sample_directory.region, TEST_USER_AGENT_PREFIX)
+    client.delete_object(
+        sample_directory.bucket,
+        f"{sample_directory.prefix}hello_world.txt",
+    )
+
+    with pytest.raises(S3Exception, match="Service error: The key does not exist"):
+        next(
+            client.get_object(
+                sample_directory.bucket, f"{sample_directory.prefix}hello_world.txt"
+            )
+        )
+
+
+def test_delete_object_does_not_exist(empty_directory):
+    client = MountpointS3Client(empty_directory.region, TEST_USER_AGENT_PREFIX)
+    client.delete_object(
+        empty_directory.bucket, f"{empty_directory.prefix}hello_world.txt"
+    )
+    # Assert no exception is thrown in case the object already doesn't exist - implicit
+
+
+def test_delete_object_invalid_bucket(empty_directory):
+    client = MountpointS3Client(empty_directory.region, TEST_USER_AGENT_PREFIX)
+    with pytest.raises(S3Exception, match="Service error: The bucket does not exist"):
+        client.delete_object(
+            f"{empty_directory.bucket}-{uuid.uuid4()}", empty_directory.prefix
+        )
+
+
 def _parse_list_result(stream: ListObjectStream, max_keys: int):
     object_infos = []
     i = 0

s3torchconnectorclient/python/tst/unit/test_mountpoint_s3_client.py

@@ -288,6 +288,33 @@ def test_mountpoint_client_pickles():
     assert client.no_sign_request == loaded.no_sign_request == expected_no_sign_request
 
 
+def test_delete_object():
+    key = "hello_world.txt"
+    mock_client = MockMountpointS3Client(REGION, MOCK_BUCKET)
+    mock_client.add_object(key, b"data")
+    client = mock_client.create_mocked_client()
+
+    client.delete_object(MOCK_BUCKET, key)
+
+    with pytest.raises(S3Exception, match="Service error: The key does not exist"):
+        client.get_object(MOCK_BUCKET, key)
+
+
+def test_delete_object_already_deleted():
+    mock_client = MockMountpointS3Client(REGION, MOCK_BUCKET)
+    client = mock_client.create_mocked_client()
+
+    client.delete_object(MOCK_BUCKET, "hello_world.txt")
+
+
+def test_delete_object_non_existent_bucket():
+    mock_client = MockMountpointS3Client(REGION, MOCK_BUCKET)
+    client = mock_client.create_mocked_client()
+
+    with pytest.raises(S3Exception, match="Service error: The bucket does not exist"):
+        client.delete_object("bucket2", "hello_world.txt")
+
+
 def _assert_isinstance(obj, expected: type):
     assert isinstance(obj, expected), f"Expected a {expected}, got {type(obj)=}"
 

s3torchconnectorclient/rust/src/mountpoint_s3_client.rs

@@ -7,8 +7,8 @@ use std::sync::Arc;
 
 use mountpoint_s3_client::config::{EndpointConfig, S3ClientAuthConfig, S3ClientConfig};
 use mountpoint_s3_client::types::PutObjectParams;
-use mountpoint_s3_client::{ObjectClient, S3CrtClient};
 use mountpoint_s3_client::user_agent::UserAgent;
+use mountpoint_s3_client::{ObjectClient, S3CrtClient};
 use nix::unistd::Pid;
 use pyo3::types::PyTuple;
 use pyo3::{pyclass, pymethods, PyRef, PyResult, ToPyObject};
@@ -64,7 +64,8 @@ impl MountpointS3Client {
         let endpoint_config = EndpointConfig::new(&region);
         let auth_config = auth_config(profile.as_deref(), no_sign_request);
 
-        let user_agent_suffix = &format!("{}/{}", build_info::PACKAGE_NAME, build_info::FULL_VERSION);
+        let user_agent_suffix =
+            &format!("{}/{}", build_info::PACKAGE_NAME, build_info::FULL_VERSION);
         let mut user_agent_string = &format!("{} {}", &user_agent_prefix, &user_agent_suffix);
         if user_agent_prefix.ends_with(user_agent_suffix) {
             // If we unpickle a client, we should not append the suffix again
@@ -130,6 +131,10 @@ impl MountpointS3Client {
         slf.client.head_object(slf.py(), bucket, key)
     }
 
+    pub fn delete_object(slf: PyRef<'_, Self>, bucket: String, key: String) -> PyResult<()> {
+        slf.client.delete_object(slf.py(), bucket, key)
+    }
+
     pub fn __getnewargs__(slf: PyRef<'_, Self>) -> PyResult<&PyTuple> {
         let py = slf.py();
         let state = [

s3torchconnectorclient/rust/src/mountpoint_s3_client_inner.rs

@@ -41,6 +41,7 @@ pub(crate) trait MountpointS3ClientInner {
         params: PutObjectParams,
     ) -> PyResult<PutObjectStream>;
     fn head_object(&self, py: Python, bucket: String, key: String) -> PyResult<PyObjectInfo>;
+    fn delete_object(&self, py: Python, bucket: String, key: String) -> PyResult<()>;
 }
 
 pub(crate) struct MountpointS3ClientInnerImpl<T: ObjectClient> {
@@ -108,4 +109,12 @@ where
         let request = py.allow_threads(|| block_on(request).map_err(python_exception))?;
         Ok(PyObjectInfo::from_object_info(request.object))
     }
+
+    fn delete_object(&self, py: Python, bucket: String, key: String) -> PyResult<()> {
+        let request = self.client.delete_object(&bucket, &key);
+
+        // TODO - Look at use of `block_on` and see if we can future this.
+        py.allow_threads(|| block_on(request).map_err(python_exception))?;
+        Ok(())
+    }
 }