From 763b654dd05076e578cc1811dd1f7fe153fd945a Mon Sep 17 00:00:00 2001
From: Joe Prosser <joe.prosser@uipath.com>
Date: Wed, 8 May 2024 13:36:16 +0100
Subject: [PATCH] feat(commands): add get buckets stats

---
 CHANGELOG.md                    |  2 ++
 cli/src/commands/get/buckets.rs | 40 +++++++++++++++++++++++++---
 cli/src/printer.rs              | 46 +++++++++++++++++++++++++++++++--
 3 files changed, 83 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9b1ba09f..ece9d1d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,8 @@
 - Fixes issue where upper case file names would not be matched in `parse` 
 - Reduce batch size when deleting comment batches
 - Support attachment type filters
+- support getting stats for `get buckets` 
+
 
 # v0.24.0
 - BREAKING: the `--context` option is now required. Users need to opt
diff --git a/cli/src/commands/get/buckets.rs b/cli/src/commands/get/buckets.rs
index 25a37686..afa1a4c9 100644
--- a/cli/src/commands/get/buckets.rs
+++ b/cli/src/commands/get/buckets.rs
@@ -1,18 +1,29 @@
+use std::collections::HashMap;
+
 use anyhow::{Context, Result};
+use log::info;
 use reinfer_client::{BucketIdentifier, Client};
 use structopt::StructOpt;
 
-use crate::printer::Printer;
+use crate::printer::{PrintableBucket, Printer};
 
 #[derive(Debug, StructOpt)]
 pub struct GetBucketsArgs {
     #[structopt(name = "bucket")]
     /// If specified, only list this bucket (name or id)
     bucket: Option<BucketIdentifier>,
+
+    #[structopt(long = "stats")]
+    /// Whether to include bucket statistics in response
+    include_stats: bool,
 }
 
 pub fn get(client: &Client, args: &GetBucketsArgs, printer: &Printer) -> Result<()> {
-    let GetBucketsArgs { bucket } = args;
+    let GetBucketsArgs {
+        bucket,
+        include_stats,
+    } = args;
+
     let buckets = if let Some(bucket) = bucket {
         vec![client
             .get_bucket(bucket.clone())
@@ -26,5 +37,28 @@ pub fn get(client: &Client, args: &GetBucketsArgs, printer: &Printer) -> Result<
         });
         buckets
     };
-    printer.print_resources(&buckets)
+
+    let mut bucket_stats: HashMap<_, _> = HashMap::new();
+
+    if *include_stats {
+        buckets.iter().try_for_each(|bucket| -> Result<()> {
+            info!("Getting statistics for bucket {}", bucket.full_name().0);
+            let stats = client
+                .get_bucket_statistics(&bucket.full_name())
+                .context("Could not get statistics for bucket")?;
+
+            bucket_stats.insert(bucket.id.clone(), stats);
+            Ok(())
+        })?;
+    }
+
+    let printable_buckets: Vec<PrintableBucket> = buckets
+        .into_iter()
+        .map(|bucket| {
+            let stats = bucket_stats.get(&bucket.id).cloned();
+            PrintableBucket { bucket, stats }
+        })
+        .collect();
+
+    printer.print_resources(&printable_buckets)
 }
diff --git a/cli/src/printer.rs b/cli/src/printer.rs
index d8cbc121..4d942c43 100644
--- a/cli/src/printer.rs
+++ b/cli/src/printer.rs
@@ -3,8 +3,8 @@ use colored::Colorize;
 use prettytable::{format, row, Row, Table};
 use reinfer_client::{
     resources::{
-        audit::PrintableAuditEvent, dataset::DatasetAndStats, integration::Integration,
-        quota::Quota,
+        audit::PrintableAuditEvent, bucket_statistics::Statistics as BucketStatistics,
+        dataset::DatasetAndStats, integration::Integration, quota::Quota,
     },
     Bucket, CommentStatistics, Dataset, Project, Source, Stream, User,
 };
@@ -182,6 +182,48 @@ impl DisplayTable for Source {
     }
 }
 
+#[derive(Debug)]
+pub struct PrintableBucket {
+    pub bucket: Bucket,
+    pub stats: Option<BucketStatistics>,
+}
+impl DisplayTable for PrintableBucket {
+    fn to_table_headers() -> Row {
+        row![bFg => "Name", "ID", "Created (UTC)", "Transform Tag", "Num Emails"]
+    }
+
+    fn to_table_row(&self) -> Row {
+        let full_name = format!(
+            "{}{}{}",
+            self.bucket.owner.0.dimmed(),
+            "/".dimmed(),
+            self.bucket.name.0
+        );
+        row![
+            full_name,
+            self.bucket.id.0,
+            self.bucket.created_at.format("%Y-%m-%d %H:%M:%S"),
+            match &self.bucket.transform_tag {
+                Some(transform_tag) => transform_tag.0.as_str().into(),
+                None => "missing".dimmed(),
+            },
+            if let Some(stats) = &self.stats {
+                stats.count.to_string().as_str().into()
+            } else {
+                "none".dimmed()
+            }
+        ]
+    }
+}
+impl Serialize for PrintableBucket {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        Serialize::serialize(&self.bucket, serializer)
+    }
+}
+
 /// Source with additional fields for printing
 /// Serializes to a Source
 #[derive(Debug)]