GH-55 # Refactor download_top_1m_file

Hugo-C · Apr 1, 2024 · 473d9b2 · 473d9b2
1 parent 8060781
commit 473d9b2
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/src/tranco_top1m/mod.rs b/src/tranco_top1m/mod.rs
@@ -82,11 +82,13 @@ impl TrancoTop1M {
             info!("Downloading tranco initialisation file");
             let tmp_file = NamedTempFile::new().unwrap();
             let path = tmp_file.path().to_str().unwrap().to_string();
-            let mut tokio_tmp_file = tokio::fs::File::from(tmp_file.into_file());
-            let mut byte_stream = reqwest::get(url).await.unwrap().bytes_stream();
 
-            while let Some(item) = byte_stream.next().await {
-                tokio::io::copy(&mut item.unwrap().as_ref(), &mut tokio_tmp_file).await.unwrap();
+            match Self::download_top_1m_file(url, tmp_file).await {
+                Ok(_) => info!("Tranco top 1M file downloaded successfully!"),
+                Err(err) => {
+                    sentry::capture_error(&err);
+                    error!("Failed to download tranco file, aborting.")
+                }
             }
             path
         };
@@ -95,6 +97,16 @@ impl TrancoTop1M {
         let _: () = self.redis_client.set(TRANCO_TOP_1M_INITIALIZED_KEY, 1).await.unwrap();
     }
 
+    async fn download_top_1m_file(url: String, file: NamedTempFile) -> Result<(), reqwest::Error>{
+        let mut tokio_tmp_file = tokio::fs::File::from(file.into_file());
+        let mut byte_stream = reqwest::get(url).await?.bytes_stream();
+
+        while let Some(item) = byte_stream.next().await {
+            tokio::io::copy(&mut item?.as_ref(), &mut tokio_tmp_file).await.unwrap();
+        }
+        Ok(())
+    }
+
     /// Remove all keys related to Tranco in the Redis DB
     pub async fn destroy_db(&mut self) -> Result<(), RedisError> {
         // First remove the init key so the other keys are not used in a partial state