feat: make new functions get_decompress and set_compress (#153)

Cargo.toml

@@ -30,6 +30,7 @@ serde_json = "1.0.79"
 thiserror = "1.0.38"
 base64 = "0.21.0"
 futures = "0"
+zstd = "0.12.4"
 
 [dev-dependencies]
 base64-url = "1.4.13"

src/compression_utils.rs

@@ -0,0 +1,62 @@
+#[inline]
+pub fn compress_json(bytes: &[u8]) -> Result<Vec<u8>, std::io::Error> {
+    zstd::encode_all(bytes, 0)
+}
+
+pub fn decompress_json(compressed: &[u8]) -> Result<Vec<u8>, std::io::Error> {
+    zstd::decode_all(compressed)
+}
+
+#[cfg(test)]
+mod test {
+
+    use serde::{Deserialize, Serialize};
+
+    use crate::compression_utils::{compress_json, decompress_json};
+
+    #[derive(Serialize, Deserialize, Default)]
+    struct JsonObject {
+        company_name: String,
+        employee_count: u32,
+        field1: String,
+        field2: u32,
+        field3: String,
+        field4: u32,
+        field5: String,
+    }
+
+    #[test]
+    fn test_compress_json() {
+        let json_object = JsonObject {
+            company_name: "momento".to_string(),
+            employee_count: 50,
+            field1: "field1".to_string(),
+            field2: 100,
+            field3: "field3".to_string(),
+            field4: 100,
+            field5: "field5".to_string(),
+        };
+        let json_str = serde_json::to_string(&json_object).expect("Serialization failed");
+        let compressed = compress_json(json_str.as_bytes()).expect("compression failed");
+
+        // assert we are in fact compressed
+        assert!(compressed.len() < json_str.len());
+        assert_eq!(
+            vec![
+                40, 181, 47, 253, 0, 88, 133, 2, 0, 66, 4, 15, 21, 160, 87, 7, 214, 33, 23, 145, 2,
+                168, 248, 92, 10, 64, 100, 88, 219, 183, 152, 254, 55, 33, 135, 17, 10, 206, 41, 6,
+                13, 16, 180, 17, 216, 137, 193, 164, 32, 180, 211, 98, 87, 250, 223, 178, 150, 55,
+                241, 58, 125, 110, 26, 253, 84, 154, 158, 170, 150, 187, 108, 4, 7, 0, 128, 16,
+                192, 15, 16, 6, 136, 16, 138, 40, 24, 236, 72, 148, 61
+            ],
+            compressed
+        );
+
+        let decompressed = decompress_json(compressed.as_slice()).expect("decompression failed");
+        assert_eq!(json_str.as_bytes(), decompressed);
+        assert_eq!(
+            json_str,
+            String::from_utf8(decompressed).expect("failed to convert from bytes to utf8")
+        )
+    }
+}

src/lib.rs

@@ -3,6 +3,7 @@ pub mod preview;
 pub mod requests;
 pub mod response;
 
+mod compression_utils;
 mod credential_provider;
 mod grpc;
 mod simple_cache_client;

src/simple_cache_client.rs

@@ -17,6 +17,7 @@ use std::ops::RangeBounds;
 use std::time::{Duration, UNIX_EPOCH};
 use tonic::{codegen::InterceptedService, transport::Channel, Request};
 
+use crate::compression_utils::{compress_json, decompress_json};
 use crate::credential_provider::CredentialProvider;
 use crate::requests::generate_api_token_request::TokenExpiry;
 use crate::response::{
@@ -529,6 +530,50 @@ impl SimpleCacheClient {
         Ok(MomentoSetResponse::new())
     }
 
+    /// Sets an item in a Momento Cache, compressing it first. Item must be retrieved with
+    /// get_with_decompression to be read properly
+    ///
+    /// # Arguments
+    ///
+    /// * `cache_name` - name of cache
+    /// * `cache_key` - key of entry within the cache.
+    /// * `cache_body` - data stored within the cache entry.
+    /// * `ttl` - The TTL to use for the
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # fn main() -> momento_test_util::DoctestResult {
+    /// # momento_test_util::doctest(|cache_name, credential_provider| async move {
+    /// use std::time::Duration;
+    /// use momento::SimpleCacheClientBuilder;
+    ///
+    /// let mut momento = SimpleCacheClientBuilder::new(credential_provider, Duration::from_secs(30))?
+    ///     .build();
+    ///
+    /// // Use client default TTL: 30 seconds, as specified above.
+    /// momento.set_with_compression(&cache_name, "k1", "v1", None).await?;
+    /// # Ok(())
+    /// # })
+    /// # }
+    /// ```
+    pub async fn set_with_compression(
+        &mut self,
+        cache_name: &str,
+        key: impl IntoBytes,
+        body: impl IntoBytes,
+        ttl: impl Into<Option<Duration>>,
+    ) -> MomentoResult<MomentoSetResponse> {
+        let compressed_body = compress_json(&body.into_bytes());
+        match compressed_body {
+            Ok(compressed) => self.set(cache_name, key, compressed, ttl).await,
+            Err(err) => Err(MomentoError::ClientSdkError {
+                description: "unable to compress json".into(),
+                source: crate::ErrorSource::Unknown(Box::new(err)),
+            }),
+        }
+    }
+
     /// Gets an item from a Momento Cache
     ///
     /// # Arguments
@@ -579,6 +624,65 @@ impl SimpleCacheClient {
         }
     }
 
+    /// Gets an item from a Momento Cache and decompresses the value before returning to the user. Item must be
+    /// set_with_compression for the return value to be correct
+    ///
+    /// # Arguments
+    ///
+    /// * `cache_name` - name of cache
+    /// * `key` - cache key
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # fn main() -> momento_test_util::DoctestResult {
+    /// # momento_test_util::doctest(|cache_name, credential_provider| async move {
+    /// use std::time::Duration;
+    /// use momento::SimpleCacheClientBuilder;
+    /// use momento::response::{Get, GetValue};
+    ///
+    /// let mut momento = SimpleCacheClientBuilder::new(credential_provider, Duration::from_secs(30))?
+    ///     .build();
+    ///
+    /// momento.set_with_compression(&cache_name, "present", "value", None).await?;
+    ///
+    /// let present = momento.get_with_decompression(&cache_name, "present").await?;
+    /// let missing = momento.get_with_decompression(&cache_name, "missing").await?;
+    ///
+    /// assert_eq!(present, Get::Hit { value: GetValue::new(b"value".to_vec()) });
+    /// assert_eq!(missing, Get::Miss);
+    /// # Ok(())
+    /// # })
+    /// # }
+    /// ```
+    pub async fn get_with_decompression(
+        &mut self,
+        cache_name: &str,
+        key: impl IntoBytes,
+    ) -> MomentoResult<Get> {
+        let get_resp = self.get(cache_name, key).await;
+        match get_resp {
+            Ok(hit) => match hit {
+                Get::Hit { value } => {
+                    let decompressed_item = decompress_json(&value.raw_item);
+                    match decompressed_item {
+                        Ok(decompressed) => Ok(Get::Hit {
+                            value: GetValue {
+                                raw_item: decompressed,
+                            },
+                        }),
+                        Err(err) => Err(MomentoError::ClientSdkError {
+                            description: "unable to decompress json".into(),
+                            source: crate::ErrorSource::Unknown(Box::new(err)),
+                        }),
+                    }
+                }
+                Get::Miss => Ok(Get::Miss),
+            },
+            Err(e) => Err(e),
+        }
+    }
+
     /// Sets dictionary items in a Momento Cache
     ///
     /// *NOTE*: This is preview functionality and requires that you contact