Add --retries parameter to CLI commands

Author: guilload

quickwit/Cargo.lock

@@ -200,10 +200,12 @@ rdkafka = { version = "0.33", default-features = false, features = [
 ] }
 regex = "1.10.0"
 regex-syntax = "0.8"
-reqwest = { version = "0.11", default-features = false, features = [
+reqwest = { version = "0.12", default-features = false, features = [
   "json",
   "rustls-tls",
 ] }
+reqwest-middleware = "0.4"
+reqwest-retry = "0.7"
 rust-embed = "6.8.1"
 rustls = "0.21"
 rustls-pemfile = "1.0.0"
@@ -304,8 +306,8 @@ azure_storage_blobs = { version = "0.13.0", default-features = false, features =
   "enable_reqwest_rustls",
 ] }
 
-opendal = { version = "0.44", default-features = false }
-reqsign = { version = "0.14", default-features = false }
+opendal = { version = "0.53", default-features = false }
+reqsign = { version = "0.16", default-features = false }
 
 quickwit-actors = { path = "quickwit-actors" }
 quickwit-aws = { path = "quickwit-aws" }
@@ -316,10 +318,10 @@ quickwit-codegen-example = { path = "quickwit-codegen/example" }
 quickwit-common = { path = "quickwit-common" }
 quickwit-config = { path = "quickwit-config" }
 quickwit-control-plane = { path = "quickwit-control-plane" }
-quickwit-index-management = { path = "quickwit-index-management" }
 quickwit-datetime = { path = "quickwit-datetime" }
 quickwit-directories = { path = "quickwit-directories" }
 quickwit-doc-mapper = { path = "quickwit-doc-mapper" }
+quickwit-index-management = { path = "quickwit-index-management" }
 quickwit-indexing = { path = "quickwit-indexing" }
 quickwit-ingest = { path = "quickwit-ingest" }
 quickwit-integration-tests = { path = "quickwit-integration-tests" }

quickwit/Cargo.toml

@@ -60,6 +60,7 @@ pub fn build_index_command() -> Command {
                         .display_order(1)
                         .required(true),
                     arg!(--overwrite "Overwrites pre-existing index. This will delete all existing data stored at `index-uri` before creating a new index.")
+                        .display_order(2)
                         .required(false),
                 ])
             )
@@ -110,6 +111,7 @@ pub fn build_index_command() -> Command {
                 .long_about("Displays descriptive statistics of an index. Displayed statistics are: number of published splits, number of documents, splits min/max timestamps, size of splits.")
                 .args(&[
                     arg!(--index <INDEX> "ID of the target index")
+                        .display_order(1)
                         .required(true),
                 ])
             )
@@ -298,7 +300,7 @@ impl IndexCliCommand {
         let client_args = ClientArgs::parse(&mut matches)?;
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
         let assume_yes = matches.get_flag("yes");
         Ok(Self::Clear(ClearIndexArgs {
             client_args,
@@ -312,7 +314,7 @@ impl IndexCliCommand {
         let index_config_uri = matches
             .remove_one::<String>("index-config")
             .map(|uri| Uri::from_str(&uri))
-            .expect("`index-config` should be a required arg.")?;
+            .expect("`index-config` should be a required arg")?;
         let overwrite = matches.get_flag("overwrite");
         let assume_yes = matches.get_flag("yes");
 
@@ -328,11 +330,11 @@ impl IndexCliCommand {
         let client_args = ClientArgs::parse(&mut matches)?;
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
         let index_config_uri = matches
             .remove_one::<String>("index-config")
             .map(|uri| Uri::from_str(&uri))
-            .expect("`index-config` should be a required arg.")?;
+            .expect("`index-config` should be a required arg")?;
         let assume_yes = matches.get_flag("yes");
 
         Ok(Self::Update(UpdateIndexArgs {
@@ -347,7 +349,8 @@ impl IndexCliCommand {
         let client_args = ClientArgs::parse(&mut matches)?;
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
+
         Ok(Self::Describe(DescribeIndexArgs {
             client_args,
             index_id,
@@ -363,7 +366,7 @@ impl IndexCliCommand {
         let client_args = ClientArgs::parse_for_ingest(&mut matches)?;
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
         let input_path_opt = if let Some(input_path) = matches.remove_one::<String>("input-path") {
             Uri::from_str(&input_path)?
                 .filepath()
@@ -401,7 +404,7 @@ impl IndexCliCommand {
     fn parse_search_args(mut matches: ArgMatches) -> anyhow::Result<Self> {
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
         let query = matches
             .remove_one::<String>("query")
             .context("`query` should be a required arg")?;
@@ -450,7 +453,7 @@ impl IndexCliCommand {
         let client_args = ClientArgs::parse(&mut matches)?;
         let index_id = matches
             .remove_one::<String>("index")
-            .expect("`index` should be a required arg.");
+            .expect("`index` should be a required arg");
         let dry_run = matches.get_flag("dry-run");
         let assume_yes = matches.get_flag("yes");
         Ok(Self::Delete(DeleteIndexArgs {

quickwit/quickwit-cli/src/index.rs

@@ -94,6 +94,16 @@ fn client_args() -> Vec<Arg> {
             .required(false)
             .global(true)
             .display_order(3),
+        Arg::new("retries")
+            .long("retries")
+            .help(
+                "Maximum number of retries for transient errors. Default value is 0. The total \
+                 number of attempts will be `1 + RETRIES`.",
+            )
+            .required(false)
+            .global(true)
+            .default_value("0")
+            .display_order(4),
     ]
 }
 
@@ -103,6 +113,7 @@ pub struct ClientArgs {
     pub connect_timeout: Option<Timeout>,
     pub timeout: Option<Timeout>,
     pub commit_timeout: Option<Timeout>,
+    pub num_retries: u32,
 }
 
 impl Default for ClientArgs {
@@ -112,6 +123,7 @@ impl Default for ClientArgs {
             connect_timeout: None,
             timeout: None,
             commit_timeout: None,
+            num_retries: 0,
         }
     }
 }
@@ -130,7 +142,7 @@ impl ClientArgs {
         if let Some(commit_timeout) = self.commit_timeout {
             builder = builder.commit_timeout(commit_timeout);
         }
-        builder
+        builder.num_retries(self.num_retries)
     }
 
     pub fn client(self) -> QuickwitClient {
@@ -149,7 +161,7 @@ impl ClientArgs {
         let cluster_endpoint = matches
             .remove_one::<String>("endpoint")
             .map(|endpoint_str| Url::from_str(&endpoint_str))
-            .expect("`endpoint` should be a required arg.")?;
+            .expect("`endpoint` should be a required arg")?;
         let connect_timeout =
             if let Some(duration) = matches.remove_one::<String>("connect-timeout") {
                 Some(parse_duration_or_none(&duration)?)
@@ -170,11 +182,16 @@ impl ClientArgs {
         } else {
             None
         };
+        let num_retries = matches
+            .remove_one::<String>("retries")
+            .map(|retries| retries.parse::<u32>())
+            .expect("`retries` should have a default value")?;
         Ok(Self {
             cluster_endpoint,
             connect_timeout,
             timeout,
             commit_timeout,
+            num_retries,
         })
     }
 }

quickwit/quickwit-cli/src/lib.rs

@@ -226,6 +226,8 @@ mod tests {
             "wikipedia",
             "--endpoint",
             "http://127.0.0.1:8000",
+            "--retries",
+            "2",
         ])?;
         let command = CliCommand::parse_cli_args(matches)?;
         assert!(matches!(
@@ -243,6 +245,7 @@ mod tests {
                 && client_args.connect_timeout.is_none()
                 && client_args.commit_timeout.is_none()
                 && client_args.cluster_endpoint == Url::from_str("http://127.0.0.1:8000").unwrap()
+                && client_args.num_retries == 2
         ));
 
         let app = build_cli().no_binary_name(true);
@@ -272,6 +275,7 @@ mod tests {
                         && client_args.timeout.is_none()
                         && client_args.connect_timeout.is_none()
                         && client_args.commit_timeout.is_none()
+                        && client_args.num_retries == 0
                         && batch_size_limit == ByteSize::mb(8)
         ));
 
@@ -301,6 +305,7 @@ mod tests {
                     && client_args.timeout.is_none()
                     && client_args.connect_timeout.is_none()
                     && client_args.commit_timeout.is_none()
+                    && client_args.num_retries == 0
                     && batch_size_limit == ByteSize::kb(4)
         ));
 
@@ -331,6 +336,7 @@ mod tests {
                         && client_args.timeout == Some(Timeout::from_secs(10))
                         && client_args.connect_timeout == Some(Timeout::from_secs(2))
                         && client_args.commit_timeout.is_none()
+                        && client_args.num_retries == 0
         ));
 
         let app = build_cli().no_binary_name(true);

quickwit/quickwit-cli/src/main.rs

@@ -13,6 +13,8 @@ license.workspace = true
 [dependencies]
 bytes = { workspace = true }
 reqwest = { workspace = true }
+reqwest-middleware = { workspace = true }
+reqwest-retry = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
@@ -30,6 +32,7 @@ quickwit-search = { workspace = true }
 quickwit-serve = { workspace = true }
 
 [dev-dependencies]
+http = { workspace = true }
 wiremock = { workspace = true }
 
 quickwit-config = { workspace = true, features = ["testsuite"] }

quickwit/quickwit-rest-client/Cargo.toml

@@ -36,6 +36,9 @@ pub enum Error {
     // Json serialization/deserialization error.
     #[error("Serde JSON error: {0}")]
     Json(#[from] serde_json::error::Error),
+    // Error returned by reqwest middleware.
+    #[error(transparent)]
+    Middleware(#[from] reqwest_middleware::Error),
     // Error returned by url lib when parsing a string.
     #[error("URL parsing error: {0}")]
     UrlParse(String),

quickwit/quickwit-rest-client/src/error.rs

@@ -27,7 +27,10 @@ use quickwit_serve::{
 };
 use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
 use reqwest::tls::Certificate;
-use reqwest::{Client, ClientBuilder, Method, StatusCode, Url};
+use reqwest::{ClientBuilder as ReqwestClientBuilder, Method, StatusCode, Url};
+use reqwest_middleware::{ClientBuilder as ReqwestMiddlewareClientBuilder, ClientWithMiddleware};
+use reqwest_retry::policies::ExponentialBackoff;
+use reqwest_retry::RetryTransientMiddleware;
 use serde::Serialize;
 use serde_json::json;
 
@@ -47,28 +50,41 @@ pub const DEFAULT_CLIENT_COMMIT_TIMEOUT: Timeout = Timeout::from_mins(30);
 struct Transport {
     base_url: Url,
     api_url: Url,
-    client: Client,
+    client: ClientWithMiddleware,
 }
 
 impl Transport {
-    fn new(endpoint: Url, connect_timeout: Timeout, ca_cert: Option<Certificate>) -> Self {
+    fn new(
+        endpoint: Url,
+        connect_timeout: Timeout,
+        ca_cert: Option<Certificate>,
+        num_retries: u32,
+    ) -> Self {
         let base_url = endpoint;
         let api_url = base_url
             .join("api/v1/")
-            .expect("Endpoint should not be malformed.");
-        let mut client_builder = ClientBuilder::new();
+            .expect("root url should be well-formed");
+        let mut reqwest_client_builder = ReqwestClientBuilder::new();
         if let Some(duration) = connect_timeout.as_duration_opt() {
-            client_builder = client_builder.connect_timeout(duration);
+            reqwest_client_builder = reqwest_client_builder.connect_timeout(duration);
         }
         if let Some(ca_cert) = ca_cert {
-            client_builder = client_builder
+            reqwest_client_builder = reqwest_client_builder
                 .tls_built_in_root_certs(false)
                 .add_root_certificate(ca_cert);
         }
+        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(num_retries);
+        let retry_transient_middleware = RetryTransientMiddleware::new_with_policy(retry_policy);
+        let reqwest_client = reqwest_client_builder
+            .build()
+            .expect("`client_builder.build()` should not fail");
+        let client = ReqwestMiddlewareClientBuilder::new(reqwest_client)
+            .with(retry_transient_middleware)
+            .build();
         Self {
             base_url,
             api_url,
-            client: client_builder.build().expect("Client should be built."),
+            client,
         }
     }
 
@@ -129,6 +145,8 @@ pub struct QuickwitClientBuilder {
     detailed_response: bool,
     /// Validate against a custom TLS certificate authority
     ca_cert: Option<Certificate>,
+    /// Maximum number of retries for transient errors.
+    num_retries: u32,
 }
 
 impl QuickwitClientBuilder {
@@ -143,6 +161,7 @@ impl QuickwitClientBuilder {
             use_legacy_ingest: false,
             detailed_response: false,
             ca_cert: None,
+            num_retries: 0,
         }
     }
 
@@ -187,8 +206,18 @@ impl QuickwitClientBuilder {
         self
     }
 
+    pub fn num_retries(mut self, num_retries: u32) -> Self {
+        self.num_retries = num_retries;
+        self
+    }
+
     pub fn build(self) -> QuickwitClient {
-        let transport = Transport::new(self.base_url, self.connect_timeout, self.ca_cert);
+        let transport = Transport::new(
+            self.base_url,
+            self.connect_timeout,
+            self.ca_cert,
+            self.num_retries,
+        );
         QuickwitClient {
             transport,
             timeout: self.timeout,
@@ -744,6 +773,7 @@ mod test {
     use std::path::PathBuf;
     use std::str::FromStr;
 
+    use http::StatusCode;
     use quickwit_config::{ConfigFormat, SourceConfig};
     use quickwit_indexing::mock_split;
     use quickwit_ingest::CommitType;
@@ -752,7 +782,7 @@ mod test {
         ListSplitsQueryParams, ListSplitsResponse, RestIngestResponse, SearchRequestQueryString,
     };
     use reqwest::header::CONTENT_TYPE;
-    use reqwest::{StatusCode, Url};
+    use reqwest::Url;
     use serde_json::json;
     use tokio::fs::File;
     use tokio::io::AsyncReadExt;
@@ -771,9 +801,7 @@ mod test {
         let server_url = Url::parse(&format!("http://127.0.0.1:{port}")).unwrap();
         let qw_client = QuickwitClientBuilder::new(server_url).build();
         let error = qw_client.indexes().list().await.unwrap_err();
-
-        assert!(matches!(error, Error::Client(_)));
-        assert!(error.to_string().contains("tcp connect error"));
+        assert!(matches!(error, Error::Middleware(_)));
     }
 
     #[tokio::test]