feat(comments): add messages filters (#248)

reinfer · Jan 9, 2024 · ecec086 · ecec086
1 parent dca3602
commit ecec086
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,6 @@
+## Unreleased
+- Add messages filters
+
 ## v0.21.3 
 - Reduce batch size for parse emls
 

diff --git a/api/src/resources/comment.rs b/api/src/resources/comment.rs
@@ -67,11 +67,11 @@ pub enum ReviewedFilterEnum {
 type UserPropertyName = String;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct UserPropertiesFilter(pub HashMap<UserPropertyName, UserPropertyFilterKind>);
+pub struct UserPropertiesFilter(pub HashMap<UserPropertyName, PropertyFilterKind>);
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
-pub enum UserPropertyFilterKind {
+pub enum PropertyFilterKind {
     OneOf(Vec<PropertyValue>),
 }
 
@@ -89,6 +89,18 @@ pub struct CommentFilter {
     #[serde(skip_serializing_if = "Vec::is_empty")]
     #[serde(default)]
     pub sources: Vec<SourceId>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub messages: Option<MessagesFilter>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct MessagesFilter {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub from: Option<PropertyFilterKind>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub to: Option<PropertyFilterKind>,
 }
 
 #[derive(Debug, Clone, Serialize)]

diff --git a/api/src/resources/dataset.rs b/api/src/resources/dataset.rs
@@ -301,7 +301,7 @@ pub(crate) struct UpdateResponse {
 #[cfg(test)]
 mod tests {
     use crate::resources::comment::{
-        CommentTimestampFilter, UserPropertiesFilter, UserPropertyFilterKind,
+        CommentTimestampFilter, PropertyFilterKind, UserPropertiesFilter,
     };
 
     use super::*;
@@ -431,7 +431,7 @@ mod tests {
     pub fn test_serialize_user_properties_request_params() {
         let user_property_filter = UserPropertiesFilter(HashMap::from([(
             "string:Generation Tag".to_string(),
-            UserPropertyFilterKind::OneOf(vec![PropertyValue::String(
+            PropertyFilterKind::OneOf(vec![PropertyValue::String(
                 "72b01fe7-ef2e-481e-934d-bc2fe0ca9b06".to_string(),
             )]),
         )]));

diff --git a/cli/src/commands/get/comments.rs b/cli/src/commands/get/comments.rs
@@ -8,8 +8,8 @@ use regex::Regex;
 use reinfer_client::{
     resources::{
         comment::{
-            CommentTimestampFilter, ReviewedFilterEnum, UserPropertiesFilter,
-            UserPropertyFilterKind,
+            CommentTimestampFilter, MessagesFilter, PropertyFilterKind, ReviewedFilterEnum,
+            UserPropertiesFilter,
         },
         dataset::{
             Attribute, AttributeFilter, AttributeFilterEnum, OrderEnum, QueryRequestParams,
@@ -89,6 +89,14 @@ pub struct GetManyCommentsArgs {
     /// Ending timestamp for comments to retrieve (inclusive).
     to_timestamp: Option<DateTime<Utc>>,
 
+    #[structopt(long = "senders")]
+    /// Filter to comments only from these senders
+    senders: Option<Vec<String>>,
+
+    #[structopt(long = "recipients")]
+    /// Filter to emails only to these recipients
+    recipients: Option<Vec<String>>,
+
     #[structopt(short = "f", long = "file", parse(from_os_str))]
     /// Path where to write comments as JSON. If not specified, stdout will be used.
     path: Option<PathBuf>,
@@ -211,10 +219,7 @@ fn get_user_properties_filter_interactively(
             }
         }
 
-        filters.insert(
-            selected_property_name,
-            UserPropertyFilterKind::OneOf(values),
-        );
+        filters.insert(selected_property_name, PropertyFilterKind::OneOf(values));
 
         if !Confirm::new()
             .with_prompt("Do you want to filter additional user properties?")
@@ -241,6 +246,8 @@ pub fn get_many(client: &Client, args: &GetManyCommentsArgs) -> Result<()> {
         label_filter,
         property_filter: user_property_filter,
         interactive_property_filter: interative_property_filter,
+        recipients,
+        senders,
     } = args;
 
     let by_timerange = from_timestamp.is_some() || to_timestamp.is_some();
@@ -287,6 +294,10 @@ pub fn get_many(client: &Client, args: &GetManyCommentsArgs) -> Result<()> {
         bail!("The `interative_property_filter` and `property_filter` options are mutually exclusive.")
     }
 
+    if (senders.is_some() || recipients.is_some()) && dataset.is_none() {
+        bail!("Cannot filter on `senders` or `recipients` when `dataset` is not provided")
+    }
+
     let file = match path {
         Some(path) => Some(
             File::create(path)
@@ -316,6 +327,25 @@ pub fn get_many(client: &Client, args: &GetManyCommentsArgs) -> Result<()> {
         None
     };
 
+    let messages_filter = MessagesFilter {
+        from: senders.as_ref().map(|senders| {
+            PropertyFilterKind::OneOf(
+                senders
+                    .iter()
+                    .map(|sender| PropertyValue::String(sender.to_owned()))
+                    .collect(),
+            )
+        }),
+        to: recipients.as_ref().map(|recipients| {
+            PropertyFilterKind::OneOf(
+                recipients
+                    .iter()
+                    .map(|recipient| PropertyValue::String(recipient.to_owned()))
+                    .collect(),
+            )
+        }),
+    };
+
     let download_options = CommentDownloadOptions {
         dataset_identifier: dataset.clone(),
         include_predictions: include_predictions.unwrap_or(false),
@@ -328,6 +358,7 @@ pub fn get_many(client: &Client, args: &GetManyCommentsArgs) -> Result<()> {
         show_progress: !no_progress,
         label_attribute_filter,
         user_properties_filter,
+        messages_filter: Some(messages_filter),
     };
 
     if let Some(file) = file {
@@ -386,6 +417,7 @@ struct CommentDownloadOptions {
     show_progress: bool,
     label_attribute_filter: Option<AttributeFilter>,
     user_properties_filter: Option<UserPropertiesFilter>,
+    messages_filter: Option<MessagesFilter>,
 }
 
 impl CommentDownloadOptions {
@@ -424,6 +456,7 @@ fn download_comments(
                 None
             },
             user_properties: options.user_properties_filter.clone(),
+            messages: options.messages_filter.clone(),
         };
 
         Ok(get_comments_progress_bar(
@@ -541,6 +574,7 @@ fn get_comments_from_uids(
             }),
             user_properties: options.user_properties_filter.clone(),
             sources: vec![source.id],
+            messages: options.messages_filter.clone(),
         },
         limit: DEFAULT_QUERY_PAGE_SIZE,
         order: OrderEnum::Recent,