From f33b1830f0c46e99cf6e49f0d5104db7ca617ab8 Mon Sep 17 00:00:00 2001
From: Alexandru-Liviu Bratosin <alexandru.bratosin@uipath.com>
Date: Fri, 28 Jun 2024 12:35:19 +0100
Subject: [PATCH] fix(stats): update api schema for bucket statistics

---
 api/src/resources/bucket_statistics.rs | 23 ++++++++++++++++++++++-
 cli/src/commands/get/emails.rs         |  9 +++++++--
 cli/src/printer.rs                     | 21 ++++++++++++++-------
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/api/src/resources/bucket_statistics.rs b/api/src/resources/bucket_statistics.rs
index acf411e..6abc89f 100644
--- a/api/src/resources/bucket_statistics.rs
+++ b/api/src/resources/bucket_statistics.rs
@@ -5,7 +5,28 @@ pub struct GetBucketStatisticsResponse {
     pub statistics: Statistics,
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+#[serde(tag = "kind")]
+pub enum Count {
+    #[serde(rename = "lower_bound")]
+    LowerBoundBucketCount { value: i32 },
+    #[serde(rename = "exact")]
+    ExactBucketCount { value: i32 },
+}
+
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 pub struct Statistics {
-    pub count: usize,
+    pub count: Count,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+pub struct LowerBoundBucketCount {
+    pub kind: String,
+    pub value: i32,
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+pub struct ExactBucketCount {
+    pub kind: String,
+    pub value: i32,
 }
diff --git a/cli/src/commands/get/emails.rs b/cli/src/commands/get/emails.rs
index ac2692c..4a9cea2 100644
--- a/cli/src/commands/get/emails.rs
+++ b/cli/src/commands/get/emails.rs
@@ -1,7 +1,7 @@
 use anyhow::{Context, Result};
 
 use colored::Colorize;
-use reinfer_client::{BucketIdentifier, Client};
+use reinfer_client::{resources::bucket_statistics::Count, BucketIdentifier, Client};
 use std::{
     fs::File,
     io::{self, BufWriter, Write},
@@ -63,7 +63,12 @@ fn download_emails(
 
     let statistics = Arc::new(Statistics::new());
 
-    let _progress = get_emails_progress_bar(bucket_statistics.count as u64, &statistics);
+    let progress_bytes = match bucket_statistics.count {
+        Count::LowerBoundBucketCount { value } => value,
+        Count::ExactBucketCount { value } => value,
+    } as u64;
+
+    let _progress = get_emails_progress_bar(progress_bytes, &statistics);
 
     client
         .get_emails_iter(&bucket.full_name(), None)
diff --git a/cli/src/printer.rs b/cli/src/printer.rs
index 3ff8a9f..d5e50a7 100644
--- a/cli/src/printer.rs
+++ b/cli/src/printer.rs
@@ -3,8 +3,11 @@ use colored::Colorize;
 use prettytable::{format, row, Row, Table};
 use reinfer_client::{
     resources::{
-        audit::PrintableAuditEvent, bucket_statistics::Statistics as BucketStatistics,
-        dataset::DatasetAndStats, integration::Integration, quota::Quota,
+        audit::PrintableAuditEvent,
+        bucket_statistics::{Count, Statistics as BucketStatistics},
+        dataset::DatasetAndStats,
+        integration::Integration,
+        quota::Quota,
     },
     Bucket, CommentStatistics, Dataset, Project, Source, Stream, User,
 };
@@ -207,15 +210,19 @@ impl DisplayTable for PrintableBucket {
             "/".dimmed(),
             self.bucket.name.0
         );
+        let count_str = if let Some(stats) = &self.stats {
+            match &stats.count {
+                Count::LowerBoundBucketCount { value } => format!(">={}", value),
+                Count::ExactBucketCount { value } => format!("={}", value),
+            }
+        } else {
+            "none".dimmed().to_string()
+        };
         row![
             full_name,
             self.bucket.id.0,
             self.bucket.created_at.format("%Y-%m-%d %H:%M:%S"),
-            if let Some(stats) = &self.stats {
-                stats.count.to_string().as_str().into()
-            } else {
-                "none".dimmed()
-            }
+            count_str
         ]
     }
 }