From 1e9d8be5976bb88c3059ff878d26814f984c10df Mon Sep 17 00:00:00 2001
From: Stanislas Polu <spolu@proton.me>
Date: Mon, 20 Nov 2023 21:15:34 +0100
Subject: [PATCH] enh: qdrant_migrator retry capabilities (#2606)

---
 core/bin/qdrant_migrator.rs | 42 ++++++++++++++++++++++++++++++++-----
 1 file changed, 37 insertions(+), 5 deletions(-)

diff --git a/core/bin/qdrant_migrator.rs b/core/bin/qdrant_migrator.rs
index 5ab4b0c5eaef..a3831e20d0fb 100644
--- a/core/bin/qdrant_migrator.rs
+++ b/core/bin/qdrant_migrator.rs
@@ -344,18 +344,34 @@ fn main() -> Result<()> {
 
                 let mut page_offset: Option<PointId> = None;
                 let mut total: usize = 0;
+                let mut retry: usize = 0;
                 loop {
                     let now = utils::now();
-                    let scroll_results = qdrant_client
+                    let scroll_results = match qdrant_client
                         .scroll(&ScrollPoints {
                             collection_name: ds.qdrant_collection(),
                             with_vectors: Some(true.into()),
                             with_payload: Some(true.into()),
                             limit: Some(points_per_request as u32),
-                            offset: page_offset,
+                            offset: page_offset.clone(),
                             ..Default::default()
                         })
-                        .await?;
+                        .await
+                    {
+                        Ok(r) => r,
+                        Err(e) => {
+                            if retry < 3 {
+                                retry += 1;
+                                utils::error(&format!(
+                                    "Error migrating points (read): retry={} error={:?}",
+                                    retry, e
+                                ));
+                                continue;
+                            } else {
+                                Err(e)?
+                            }
+                        }
+                    };
 
                     let count = scroll_results.result.len();
 
@@ -371,9 +387,24 @@ fn main() -> Result<()> {
                         })
                         .collect::<Vec<_>>();
 
-                    shadow_write_qdrant_client
+                    match shadow_write_qdrant_client
                         .upsert_points(ds.qdrant_collection(), points, None)
-                        .await?;
+                        .await
+                    {
+                        Ok(_) => (),
+                        Err(e) => {
+                            if retry < 3 {
+                                retry += 1;
+                                utils::error(&format!(
+                                    "Error migrating points (write): retry={} error={:?}",
+                                    retry, e
+                                ));
+                                continue;
+                            } else {
+                                Err(e)?
+                            }
+                        }
+                    }
 
                     total += count;
                     utils::info(&format!(
@@ -387,6 +418,7 @@ fn main() -> Result<()> {
                     if page_offset.is_none() {
                         break;
                     }
+                    retry = 0;
                 }
 
                 utils::info(&format!("Done migrating: total={}", total));