From 62ab5c19e21a6905f4532737c5e1a3b743a7633e Mon Sep 17 00:00:00 2001
From: Joe Prosser <joe.prosser@uipath.com>
Date: Wed, 24 Jul 2024 14:07:19 +0100
Subject: [PATCH] feat(commands): add lossy flag when creating annotations

---
 cli/src/commands/create/annotations.rs | 23 +++++++++++++++++++++--
 cli/src/commands/create/comments.rs    |  9 +++++++++
 cli/src/main.rs                        | 19 +++++++++++++++----
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/cli/src/commands/create/annotations.rs b/cli/src/commands/create/annotations.rs
index 76722307..441761e4 100644
--- a/cli/src/commands/create/annotations.rs
+++ b/cli/src/commands/create/annotations.rs
@@ -1,4 +1,7 @@
-use crate::progress::{Options as ProgressOptions, Progress};
+use crate::{
+    print_error_as_warning,
+    progress::{Options as ProgressOptions, Progress},
+};
 use anyhow::{Context, Result};
 use colored::Colorize;
 use log::info;
@@ -47,6 +50,10 @@ pub struct CreateAnnotationsArgs {
     #[structopt(long = "batch-size", default_value = "128")]
     /// Number of comments to batch in a single request.
     batch_size: usize,
+
+    #[structopt(long)]
+    /// Whether to attempt to resume processing on error
+    lossy: bool,
 }
 
 pub fn create(client: &Client, args: &CreateAnnotationsArgs, pool: &mut Pool) -> Result<()> {
@@ -95,6 +102,7 @@ pub fn create(client: &Client, args: &CreateAnnotationsArgs, pool: &mut Pool) ->
                 args.use_moon_forms,
                 args.batch_size,
                 pool,
+                args.lossy,
             )?;
             if let Some(mut progress) = progress {
                 progress.done();
@@ -117,6 +125,7 @@ pub fn create(client: &Client, args: &CreateAnnotationsArgs, pool: &mut Pool) ->
                 args.use_moon_forms,
                 args.batch_size,
                 pool,
+                args.lossy,
             )?;
             statistics
         }
@@ -133,6 +142,7 @@ pub trait AnnotationStatistic {
     fn add_annotation(&self);
 }
 
+#[allow(clippy::too_many_arguments)]
 pub fn upload_batch_of_annotations(
     annotations_to_upload: &mut Vec<NewAnnotation>,
     client: &Client,
@@ -141,6 +151,7 @@ pub fn upload_batch_of_annotations(
     dataset_name: &DatasetFullName,
     use_moon_forms: bool,
     pool: &mut Pool,
+    lossy: bool,
 ) -> Result<()> {
     let (error_sender, error_receiver) = channel();
 
@@ -190,7 +201,12 @@ pub fn upload_batch_of_annotations(
     });
 
     if let Ok(error) = error_receiver.try_recv() {
-        Err(error)
+        if lossy {
+            print_error_as_warning(&error);
+            Ok(())
+        } else {
+            Err(error)
+        }
     } else {
         annotations_to_upload.clear();
         Ok(())
@@ -207,6 +223,7 @@ fn upload_annotations_from_reader(
     use_moon_forms: bool,
     batch_size: usize,
     pool: &mut Pool,
+    lossy: bool,
 ) -> Result<()> {
     let mut annotations_to_upload = Vec::new();
 
@@ -224,6 +241,7 @@ fn upload_annotations_from_reader(
                     dataset_name,
                     use_moon_forms,
                     pool,
+                    lossy,
                 )?;
             }
         }
@@ -238,6 +256,7 @@ fn upload_annotations_from_reader(
             dataset_name,
             use_moon_forms,
             pool,
+            lossy,
         )?;
     }
 
diff --git a/cli/src/commands/create/comments.rs b/cli/src/commands/create/comments.rs
index e28624f1..958237ab 100644
--- a/cli/src/commands/create/comments.rs
+++ b/cli/src/commands/create/comments.rs
@@ -74,6 +74,10 @@ pub struct CreateCommentsArgs {
     #[structopt(short = "y", long = "yes")]
     /// Consent to ai unit charge. Suppresses confirmation prompt.
     yes: bool,
+
+    #[structopt(long)]
+    /// Whether to attempt to resume processing on error
+    lossy: bool,
 }
 
 pub fn create(client: &Client, args: &CreateCommentsArgs, pool: &mut Pool) -> Result<()> {
@@ -152,6 +156,7 @@ pub fn create(client: &Client, args: &CreateCommentsArgs, pool: &mut Pool) -> Re
                 args.use_moon_forms,
                 args.no_charge,
                 pool,
+                args.lossy,
             )?;
             if let Some(mut progress) = progress {
                 progress.done();
@@ -180,6 +185,7 @@ pub fn create(client: &Client, args: &CreateCommentsArgs, pool: &mut Pool) -> Re
                 args.use_moon_forms,
                 args.no_charge,
                 pool,
+                args.lossy,
             )?;
             statistics
         }
@@ -330,6 +336,7 @@ fn upload_comments_from_reader(
     use_moon_forms: bool,
     no_charge: bool,
     pool: &mut Pool,
+    lossy: bool,
 ) -> Result<()> {
     assert!(batch_size > 0);
 
@@ -415,6 +422,7 @@ fn upload_comments_from_reader(
                     dataset_name,
                     use_moon_forms,
                     pool,
+                    lossy,
                 )?;
             }
         }
@@ -442,6 +450,7 @@ fn upload_comments_from_reader(
                 dataset_name,
                 use_moon_forms,
                 pool,
+                lossy,
             )?;
         }
     }
diff --git a/cli/src/main.rs b/cli/src/main.rs
index 3d4e79f4..d504e5c2 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -203,15 +203,26 @@ fn find_configuration(args: &Args) -> Result<PathBuf> {
     Ok(config_path)
 }
 
+pub fn print_error_as_warning(error: &anyhow::Error) {
+    warn!("An error occurred. Resuming...:");
+    for cause in error.chain() {
+        warn!(" |- {}", cause);
+    }
+}
+
+pub fn print_error(error: &anyhow::Error) {
+    error!("An error occurred:");
+    for cause in error.chain() {
+        error!(" |- {}", cause);
+    }
+}
+
 fn main() {
     let args = Args::from_args();
     utils::init_env_logger(args.verbose);
 
     if let Err(error) = run(args) {
-        error!("An error occurred:");
-        for cause in error.chain() {
-            error!(" |- {}", cause);
-        }
+        print_error(&error);
 
         #[cfg(feature = "backtrace")]
         {