From b2a6babad0b06df3a46c1efe7dc299bfe777c6b2 Mon Sep 17 00:00:00 2001 From: Matt Straathof <11823378+bruuuuuuuce@users.noreply.github.com> Date: Tue, 19 Sep 2023 15:21:06 -0700 Subject: [PATCH] feat: make new functions get_decompress and set_compress (#153) --- Cargo.toml | 1 + src/compression_utils.rs | 62 ++++++++++++++++++++++ src/lib.rs | 1 + src/simple_cache_client.rs | 104 +++++++++++++++++++++++++++++++++++++ 4 files changed, 168 insertions(+) create mode 100644 src/compression_utils.rs diff --git a/Cargo.toml b/Cargo.toml index 13e3ca15..e55633ef 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/compression_utils.rs b/src/compression_utils.rs new file mode 100644 index 00000000..361d3979 --- /dev/null +++ b/src/compression_utils.rs @@ -0,0 +1,62 @@ +#[inline] +pub fn compress_json(bytes: &[u8]) -> Result, std::io::Error> { + zstd::encode_all(bytes, 0) +} + +pub fn decompress_json(compressed: &[u8]) -> Result, 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") + ) + } +} diff --git a/src/lib.rs b/src/lib.rs index 3d314913..d318c3e3 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/simple_cache_client.rs b/src/simple_cache_client.rs index a014d45e..735c0418 100644 --- a/src/simple_cache_client.rs +++ b/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>, + ) -> MomentoResult { + 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 { + 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