refactor: Refactored the download code so that it passes all linter checks

This commit also prepares the code to be used as a library and to make bindings from it.

Author: pjox

src/download.rs

@@ -2,27 +2,64 @@ use flate2::read::GzDecoder;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use reqwest::{header, Client, Url};
 use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
-use reqwest_retry::policies::ExponentialBackoff;
-use reqwest_retry::Jitter;
-use reqwest_retry::RetryTransientMiddleware;
-use std::fs::File;
-use std::io::BufRead;
-use std::io::BufReader;
-use std::path::{Path, PathBuf};
-use std::process;
-use std::sync::Arc;
-use std::time::Duration;
-use tokio::io::AsyncWriteExt;
-use tokio::io::BufWriter;
-use tokio::sync::Semaphore;
-use tokio::task::JoinSet;
+use reqwest_retry::{policies::ExponentialBackoff, Jitter, RetryTransientMiddleware};
+use std::{
+    fs::File,
+    io::{BufRead, BufReader},
+    path::{Path, PathBuf},
+    process, str,
+    sync::Arc,
+    time::Duration,
+};
+use tokio::{
+    io::{AsyncWriteExt, BufWriter},
+    sync::Semaphore,
+    task::JoinSet,
+};
 
 use crate::errors::DownloadError;
 
 const BASE_URL: &str = "https://data.commoncrawl.org/";
 
 static APP_USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"),);
 
+pub struct DownloadOptions<'a> {
+    pub snapshot: String,
+    pub data_type: &'a str,
+    pub paths: &'a Path,
+    pub dst: &'a Path,
+    pub threads: usize,
+    pub max_retries: usize,
+    pub numbered: bool,
+    pub files_only: bool,
+    pub progress: bool,
+}
+
+pub struct TaskOptions {
+    pub number: usize,
+    pub path: String,
+    pub dst: PathBuf,
+    pub numbered: bool,
+    pub files_only: bool,
+    pub progress: bool,
+}
+
+impl Default for DownloadOptions<'_> {
+    fn default() -> Self {
+        DownloadOptions {
+            snapshot: "".to_string(),
+            data_type: "",
+            paths: Path::new(""),
+            dst: Path::new(""),
+            threads: 1,
+            max_retries: 1000,
+            numbered: false,
+            files_only: false,
+            progress: false,
+        }
+    }
+}
+
 fn new_client(max_retries: usize) -> Result<ClientWithMiddleware, DownloadError> {
     let retry_policy = ExponentialBackoff::builder()
         .retry_bounds(Duration::from_secs(1), Duration::from_secs(3600))
@@ -37,16 +74,15 @@ fn new_client(max_retries: usize) -> Result<ClientWithMiddleware, DownloadError>
         .build())
 }
 
-pub async fn download_paths(
-    snapshot: &String,
-    data_type: &str,
-    dst: &Path,
-) -> Result<(), DownloadError> {
-    let paths = format!("{}crawl-data/{}/{}.paths.gz", BASE_URL, snapshot, data_type);
+pub async fn download_paths(options: DownloadOptions<'_>) -> Result<(), DownloadError> {
+    let paths = format!(
+        "{}crawl-data/{}/{}.paths.gz",
+        BASE_URL, options.snapshot, options.data_type
+    );
     println!("Downloading paths from: {}", paths);
     let url = Url::parse(&paths)?;
 
-    let client = new_client(1000)?;
+    let client = new_client(options.max_retries)?;
 
     let filename = url
         .path_segments() // Splits into segments of the URL
@@ -55,7 +91,7 @@ pub async fn download_paths(
 
     let request = client.get(url.as_str());
 
-    let mut dst = dst.to_path_buf();
+    let mut dst = options.dst.to_path_buf();
 
     dst.push(filename);
 
@@ -79,16 +115,11 @@ pub async fn download_paths(
 
 async fn download_task(
     client: ClientWithMiddleware,
-    download_url: String,
-    number: usize,
     multibar: Arc<MultiProgress>,
-    dst: PathBuf,
-    numbered: bool,
-    files_only: bool,
-    progress: bool,
+    task_options: TaskOptions,
 ) -> Result<(), DownloadError> {
     // Parse URL into Url type
-    let url = Url::parse(&download_url)?;
+    let url = Url::parse(&task_options.path)?;
 
     // We need to determine the file size before we download, so we can create a ProgressBar
     // A Header request for the CONTENT_LENGTH header gets us the file size
@@ -109,17 +140,17 @@ async fn download_task(
     };
 
     // Parse the filename from the given URL
-    let filename = if numbered {
-        &format!("{}{}", number, ".txt.gz")
-    } else if files_only {
+    let filename = if task_options.numbered {
+        &format!("{}{}", task_options.number, ".txt.gz")
+    } else if task_options.files_only {
         url.path_segments()
             .and_then(|segments| segments.last())
             .unwrap_or("file.download")
     } else {
         url.path().strip_prefix("/").unwrap_or("file.download")
     };
 
-    let mut dst = dst.clone();
+    let mut dst = task_options.dst.clone();
 
     dst.push(filename);
 
@@ -130,7 +161,7 @@ async fn download_task(
     // and add it to the multibar
     let progress_bar = multibar.add(ProgressBar::new(download_size));
 
-    if progress {
+    if task_options.progress {
         // Set Style to the ProgressBar
         progress_bar.set_style(
             ProgressStyle::default_bar()
@@ -145,7 +176,7 @@ async fn download_task(
     }
 
     // Create the directory if it doesn't exist
-    if !numbered {
+    if !task_options.numbered {
         if let Some(parent) = dst.parent() {
             tokio::fs::create_dir_all(parent).await?;
         }
@@ -163,13 +194,13 @@ async fn download_task(
     // We use the part from the reqwest-tokio example here on purpose
     // This way, we are able to increase the ProgressBar with every downloaded chunk
     while let Some(chunk) = download.chunk().await? {
-        if progress {
+        if task_options.progress {
             progress_bar.inc(chunk.len() as u64); // Increase ProgressBar by chunk size
         }
         outfile.write_all(&chunk).await?; // Write chunk to output file
     }
 
-    if progress {
+    if task_options.progress {
         // Finish the progress bar to prevent glitches
         progress_bar.finish();
 
@@ -187,22 +218,18 @@ async fn download_task(
     Ok(())
 }
 
-pub async fn download(
-    paths: &Path,
-    dst: &Path,
-    threads: usize,
-    max_retries: usize,
-    numbered: bool,
-    files_only: bool,
-    progress: bool,
-) -> Result<(), DownloadError> {
+pub async fn download(options: DownloadOptions<'_>) -> Result<(), DownloadError> {
     // A vector containing all the URLs to download
 
     let file = {
-        let gzip_file = match File::open(paths) {
+        let gzip_file = match File::open(options.paths) {
             Ok(file) => file,
             Err(e) => {
-                eprintln!("Could not open file {}\nError: {}", paths.display(), e);
+                eprintln!(
+                    "Could not open file {}\nError: {}",
+                    options.paths.display(),
+                    e
+                );
                 process::exit(1)
             }
         };
@@ -232,7 +259,7 @@ pub async fn download(
     );
 
     // Only set the style if we are showing progress
-    if progress {
+    if options.progress {
         main_pb.set_style(
             indicatif::ProgressStyle::default_bar().template("{msg} {bar:10} {pos}/{len}")?,
         );
@@ -243,25 +270,30 @@ pub async fn download(
         main_pb.tick();
     }
 
-    let client = new_client(max_retries)?;
+    let client = new_client(options.max_retries)?;
 
-    let semaphore = Arc::new(Semaphore::new(threads));
+    let semaphore = Arc::new(Semaphore::new(options.threads));
     let mut set = JoinSet::new();
 
     for (number, path) in paths {
         // Clone multibar and main_pb.  We will move the clones into each task.
         let multibar = multibar.clone();
         let main_pb = main_pb.clone();
         let client = client.clone();
-        let dst = dst.to_path_buf();
+        let dst = options.dst.to_path_buf();
         let semaphore = semaphore.clone();
         set.spawn(async move {
             let _permit = semaphore.acquire().await;
-            let res = download_task(
-                client, path, number, multibar, dst, numbered, files_only, progress,
-            )
-            .await;
-            if progress {
+            let task_options = TaskOptions {
+                path,
+                number,
+                dst,
+                numbered: options.numbered,
+                files_only: options.files_only,
+                progress: options.progress,
+            };
+            let res = download_task(client, multibar, task_options).await;
+            if options.progress {
                 // Increment the main progress bar.
                 main_pb.inc(1);
             }
@@ -289,7 +321,7 @@ pub async fn download(
         }
     }
 
-    if progress {
+    if options.progress {
         // Change the message on the overall progress indicator.
         main_pb.finish_with_message("done");
 

src/main.rs

@@ -16,7 +16,13 @@ async fn main() {
             data_type,
             dst,
         }) => {
-            match download::download_paths(snapshot, data_type.as_str(), dst).await {
+            let options = download::DownloadOptions {
+                snapshot: snapshot.to_string(),
+                data_type: data_type.as_str(),
+                dst,
+                ..Default::default()
+            };
+            match download::download_paths(options).await {
                 Ok(_) => (),
                 Err(e) => {
                     eprintln!("Error downloading paths: {}", e);
@@ -35,17 +41,17 @@ async fn main() {
             if *numbered && *files_only {
                 eprintln!("Numbered and Files Only flags are incompatible");
             } else {
-                match download::download(
-                    path_file,
+                let options = download::DownloadOptions {
+                    paths: path_file,
                     dst,
-                    *threads,
-                    *retries,
-                    *numbered,
-                    *files_only,
-                    *progress,
-                )
-                .await
-                {
+                    progress: *progress,
+                    threads: *threads,
+                    max_retries: *retries,
+                    numbered: *numbered,
+                    files_only: *files_only,
+                    ..Default::default()
+                };
+                match download::download(options).await {
                     Ok(_) => (),
                     Err(e) => {
                         eprintln!("Error downloading paths: {}", e);