From 6663ecf6b93567d4df8a389ce9bdcd503efc76c2 Mon Sep 17 00:00:00 2001
From: Joe Prosser <joe.prosser@uipath.com>
Date: Mon, 9 Oct 2023 23:05:42 +0100
Subject: [PATCH] commands: add get dataset stats

---
 CHANGELOG.md                     |  1 +
 api/src/resources/dataset.rs     | 11 +++----
 cli/src/commands/get/datasets.rs | 49 ++++++++++++++++++++++++++++++--
 cli/src/printer.rs               |  9 ++++--
 4 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b1688053..af5ebb56 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## Unreleased
 
+- Add ability to get dataset stats
 - Show Global Permissions in `get users`
 - Upgrade `ordered-float` version, which is exposed in the public crate api.
 
diff --git a/api/src/resources/dataset.rs b/api/src/resources/dataset.rs
index b06c570b..3f366a7c 100644
--- a/api/src/resources/dataset.rs
+++ b/api/src/resources/dataset.rs
@@ -17,7 +17,7 @@ use std::{
     str::FromStr,
 };
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub struct Dataset {
     pub id: Id,
     pub name: Name,
@@ -34,6 +34,7 @@ pub struct Dataset {
     pub entity_defs: Vec<EntityDef>,
     pub label_defs: Vec<LabelDef>,
     pub label_groups: Vec<LabelGroup>,
+    pub num_reviewed: Option<f64>,
 }
 
 impl Dataset {
@@ -248,17 +249,17 @@ pub(crate) struct CreateRequest<'request> {
     pub dataset: NewDataset<'request>,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub(crate) struct CreateResponse {
     pub dataset: Dataset,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub(crate) struct GetAvailableResponse {
     pub datasets: Vec<Dataset>,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub(crate) struct GetResponse {
     pub dataset: Dataset,
 }
@@ -280,7 +281,7 @@ pub(crate) struct UpdateRequest<'request> {
     pub dataset: UpdateDataset<'request>,
 }
 
-#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub(crate) struct UpdateResponse {
     pub dataset: Dataset,
 }
diff --git a/cli/src/commands/get/datasets.rs b/cli/src/commands/get/datasets.rs
index 13f5b0ef..7c649d86 100644
--- a/cli/src/commands/get/datasets.rs
+++ b/cli/src/commands/get/datasets.rs
@@ -1,5 +1,10 @@
+use std::collections::HashMap;
+
 use anyhow::{Context, Result};
-use reinfer_client::{Client, DatasetIdentifier};
+use log::info;
+use reinfer_client::{
+    resources::dataset::StatisticsRequestParams, Client, CommentFilter, DatasetIdentifier,
+};
 use structopt::StructOpt;
 
 use crate::printer::Printer;
@@ -9,11 +14,18 @@ pub struct GetDatasetsArgs {
     #[structopt(name = "dataset")]
     /// If specified, only list this dataset (name or id)
     dataset: Option<DatasetIdentifier>,
+
+    #[structopt(long = "stats")]
+    /// Whether to include source statistics in response
+    include_stats: bool,
 }
 
 pub fn get(client: &Client, args: &GetDatasetsArgs, printer: &Printer) -> Result<()> {
-    let GetDatasetsArgs { dataset } = args;
-    let datasets = if let Some(dataset) = dataset {
+    let GetDatasetsArgs {
+        dataset,
+        include_stats,
+    } = args;
+    let mut datasets = if let Some(dataset) = dataset {
         vec![client
             .get_dataset(dataset.clone())
             .context("Operation to list datasets has failed.")?]
@@ -26,5 +38,36 @@ pub fn get(client: &Client, args: &GetDatasetsArgs, printer: &Printer) -> Result
         });
         datasets
     };
+
+    let mut dataset_stats: HashMap<_, _> = HashMap::new();
+    if *include_stats {
+        datasets.iter().try_for_each(|dataset| -> Result<()> {
+            info!("Getting statistics for dataset {}", dataset.full_name().0);
+            let stats = client
+                .get_dataset_statistics(
+                    &dataset.full_name(),
+                    &StatisticsRequestParams {
+                        comment_filter: CommentFilter {
+                            reviewed:Some(reinfer_client::resources::comment::ReviewedFilterEnum::OnlyReviewed),
+                            ..Default::default()
+                        },
+                        ..Default::default()
+                    },
+                )
+                .context("Could not get statistics for dataset")?;
+
+            dataset_stats.insert(dataset.id.clone(), stats);
+            Ok(())
+        })?;
+
+        datasets.iter_mut().for_each(|dataset| {
+            dataset.num_reviewed = if let Some(stats) = dataset_stats.get(&dataset.id).cloned() {
+                Some(*stats.num_comments)
+            } else {
+                None
+            }
+        });
+    }
+
     printer.print_resources(&datasets)
 }
diff --git a/cli/src/printer.rs b/cli/src/printer.rs
index 375f49e5..f6f7c67f 100644
--- a/cli/src/printer.rs
+++ b/cli/src/printer.rs
@@ -89,7 +89,7 @@ impl DisplayTable for Quota {
 
 impl DisplayTable for Dataset {
     fn to_table_headers() -> Row {
-        row![bFg => "Name", "ID", "Updated (UTC)", "Title"]
+        row![bFg => "Name", "ID", "Updated (UTC)", "Title", "Num Reviewed"]
     }
 
     fn to_table_row(&self) -> Row {
@@ -98,7 +98,12 @@ impl DisplayTable for Dataset {
             full_name,
             self.id.0,
             self.updated_at.format("%Y-%m-%d %H:%M:%S"),
-            self.title
+            self.title,
+            if let Some(num_reviewed) = &self.num_reviewed {
+                num_reviewed.to_string().as_str().into()
+            } else {
+                "none".dimmed()
+            }
         ]
     }
 }