From d8f552f738d5c28436da79fdf751e76ff2ff4907 Mon Sep 17 00:00:00 2001
From: Lukas <notlukas555@gmail.com>
Date: Wed, 30 Oct 2024 15:47:48 +0000
Subject: [PATCH] Sessions: add generic setters to client, add generic remove
 to client and server (#358)

---
 libs/pavex_session/src/session_.rs | 34 ++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/libs/pavex_session/src/session_.rs b/libs/pavex_session/src/session_.rs
index 28971610f..fefbbfe74 100644
--- a/libs/pavex_session/src/session_.rs
+++ b/libs/pavex_session/src/session_.rs
@@ -270,6 +270,19 @@ impl<'session> ClientSessionState<'session> {
 pub struct ClientSessionStateMut<'session>(&'session mut ClientState);
 
 impl<'session> ClientSessionStateMut<'session> {
+    /// Set a value in the client-side state for the given key.
+    ///
+    /// If the key already exists, the value is updated and the old raw value is returned.
+    /// If the value cannot be serialized, an error is returned.
+    pub fn set<T: Serialize>(
+        &mut self,
+        key: String,
+        value: T
+    ) -> Result<Option<Value>, serde_json::Error> {
+        let value = serde_json::to_value(value)?;
+        Ok(self.set_value(key, value))
+    }
+
     /// Set a value in the client-side state for the given key.
     ///
     /// If the key already exists, the value is updated and the old value is returned.
@@ -292,6 +305,16 @@ impl<'session> ClientSessionStateMut<'session> {
         }
     }
 
+    /// Remove the value associated with `key` from the client-side state.
+    ///
+    /// If the key exists, the removed value is returned.
+    /// If the removed value cannot be serialized, an error is returned.
+    pub fn remove<T: DeserializeOwned>(&mut self, key: &str) -> Result<Option<T>, serde_json::Error> {
+        self.remove_value(key)
+            .map(|value| serde_json::from_value(value))
+            .transpose()
+    }
+
     /// Remove the value associated with `key` from the client-side state.
     ///
     /// If the key exists, the removed value is returned.
@@ -419,6 +442,17 @@ impl<'session, 'store> ServerSessionState<'session, 'store> {
         Ok(old_value)
     }
 
+    /// Remove the value associated with `key` from the server-side state.
+    ///
+    /// If the key exists, the removed value is returned.
+    /// If the removed value cannot be serialized, an error is returned.
+    pub async fn remove<T: DeserializeOwned>(&mut self, key: &str) -> Result<Option<T>, LoadError> {
+        self.remove_value(key).await?
+            .map(serde_json::from_value)
+            .transpose()
+            .map_err(LoadError::DeserializationError)
+    }
+
     /// Remove the value associated with `key` from the server-side state.
     ///
     /// If the key exists, the removed value is returned.