iota_data_ingestion_core/util.rs
1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{str::FromStr, time::Duration};
6
7use backoff::{ExponentialBackoff, backoff::Backoff};
8use object_store::{
9 ClientOptions, ObjectStore, RetryConfig, aws::AmazonS3ConfigKey, gcp::GoogleConfigKey,
10};
11use url::Url;
12
13use crate::IngestionResult;
14
15/// Creates a remote store client *without* any retry mechanism.
16///
17/// This function constructs a remote store client configured to *not* retry
18/// failed requests. All requests will fail immediately if the underlying
19/// operation encounters an error. This is a convenience wrapper around
20/// `create_remote_store_client_with_ops` that sets the retry configuration
21/// to disable retries.
22///
23/// # Arguments
24///
25/// * `url`: The URL of the remote store. The scheme of the URL determines the storage provider:
26/// * `http://` or `https://`: HTTP-based store.
27/// * `gs://`: Google Cloud Storage.
28/// * `s3://` or other AWS S3-compatible URL: Amazon S3.
29/// * `remote_store_options`: A vector of key-value pairs representing provider-specific options.
30/// * For GCS: See [`object_store::gcp::GoogleConfigKey`] for valid keys.
31/// * For S3: See [`object_store::aws::AmazonS3ConfigKey`] for valid keys.
32/// * For HTTP: No options are currently supported. This parameter should be empty.
33/// * `request_timeout_secs`: The timeout duration (in seconds) for individual requests. This
34/// timeout is used to set a slightly longer retry timeout (request_timeout_secs + 1) internally,
35/// even though retries are disabled. This is done to ensure that the overall operation doesn't
36/// hang indefinitely.
37///
38/// # Examples
39///
40/// Creating an S3 client without retries:
41///
42/// ```rust,no_run
43/// # use iota_data_ingestion_core::create_remote_store_client;
44/// use object_store::aws::AmazonS3ConfigKey;
45///
46/// let url = "s3://my-bucket";
47/// let options = vec![(
48/// AmazonS3ConfigKey::Region.as_ref().to_owned(),
49/// "us-east-1".to_string(),
50/// )];
51/// let client = create_remote_store_client(url.to_string(), options, 30).unwrap();
52/// ```
53///
54/// Creating a GCS client without retries:
55///
56/// ```rust,no_run
57/// # use iota_data_ingestion_core::create_remote_store_client;
58/// use object_store::gcp::GoogleConfigKey;
59///
60/// let url = "gs://my-bucket";
61/// let options = vec![(
62/// GoogleConfigKey::ServiceAccount.as_ref().to_owned(),
63/// "my-service-account".to_string(),
64/// )];
65/// let client = create_remote_store_client(url.to_string(), options, 30).unwrap();
66/// ```
67///
68/// Creating an HTTP client without retries (no options supported):
69///
70/// ```rust,no_run
71/// # use iota_data_ingestion_core::create_remote_store_client;
72///
73/// let url = "http://example.bucket.com";
74/// let options = vec![]; // No options for HTTP
75/// let client = create_remote_store_client(url.to_string(), options, 30).unwrap();
76/// ```
77pub fn create_remote_store_client(
78 url: String,
79 remote_store_options: Vec<(String, String)>,
80 request_timeout_secs: u64,
81) -> IngestionResult<Box<dyn ObjectStore>> {
82 let retry_config = RetryConfig {
83 max_retries: 0,
84 retry_timeout: Duration::from_secs(request_timeout_secs + 1),
85 ..Default::default()
86 };
87
88 create_remote_store_client_with_ops(
89 url,
90 remote_store_options,
91 request_timeout_secs,
92 retry_config,
93 )
94}
95
96/// Creates a remote store client with configurable retry behavior and options.
97///
98/// This function constructs a remote store client for various cloud storage
99/// providers (HTTP, Google Cloud Storage, Amazon S3) based on the provided URL
100/// and options. It allows configuring retry behavior through the `retry_config`
101/// argument.
102///
103/// # Arguments
104///
105/// * `url`: The URL of the remote store. The scheme of the URL determines the storage provider:
106/// * `http://` or `https://`: HTTP-based store.
107/// * `gs://`: Google Cloud Storage.
108/// * `s3://` or other AWS S3-compatible URL: Amazon S3.
109/// * `remote_store_options`: A vector of key-value pairs representing provider-specific options.
110/// * For GCS: See [`object_store::gcp::GoogleConfigKey`] for valid keys.
111/// * For S3: See [`object_store::aws::AmazonS3ConfigKey`] for valid keys.
112/// * For HTTP: No options are currently supported. This parameter should be empty.
113/// * `request_timeout_secs`: The timeout duration (in seconds) for individual requests.
114/// * `retry_config`: A [`RetryConfig`] struct defining the retry strategy. This allows fine-grained
115/// control over the number of retries, backoff behavior, and retry timeouts. See the
116/// documentation for [`object_store::RetryConfig`] for details.
117///
118/// # Examples
119///
120/// Creating an S3 client with specific options and a retry configuration:
121///
122/// ```text
123/// # use iota_data_ingestion_core::create_remote_store_client_with_ops;
124/// use object_store::{RetryConfig, aws::AmazonS3ConfigKey};
125///
126/// let url = "s3://my-bucket";
127/// let options = vec![(
128/// AmazonS3ConfigKey::Region.as_ref().to_owned(),
129/// "us-east-1".to_string(),
130/// )];
131/// let retry_config = RetryConfig::default(); // Use default retry settings
132/// let client =
133/// create_remote_store_client_with_ops(url.to_string(), options, 30, retry_config).unwrap();
134/// ```
135///
136/// Creating a GCS client:
137///
138/// ```text
139/// # use iota_data_ingestion_core::create_remote_store_client_with_ops;
140/// use object_store::{RetryConfig, gcp::GoogleConfigKey};
141///
142/// let url = "gs://my-bucket";
143/// let options = vec![(
144/// GoogleConfigKey::ServiceAccount.as_ref().to_owned(),
145/// "my-service-account".to_string(),
146/// )];
147/// let retry_config = RetryConfig::default();
148/// let client =
149/// create_remote_store_client_with_ops(url.to_string(), options, 30, retry_config).unwrap();
150/// ```
151///
152/// Creating an HTTP client (no options supported):
153///
154/// ```text
155/// # use iota_data_ingestion_core::create_remote_store_client_with_ops;
156/// use object_store::RetryConfig;
157///
158/// let url = "http://example.bucket.com";
159/// let options = vec![]; // No options for HTTP
160/// let retry_config = RetryConfig::default();
161/// let client =
162/// create_remote_store_client_with_ops(url.to_string(), options, 30, retry_config).unwrap();
163/// ```
164pub fn create_remote_store_client_with_ops(
165 url: String,
166 remote_store_options: Vec<(String, String)>,
167 request_timeout_secs: u64,
168 retry_config: RetryConfig,
169) -> IngestionResult<Box<dyn ObjectStore>> {
170 let client_options = ClientOptions::new()
171 .with_timeout(Duration::from_secs(request_timeout_secs))
172 .with_allow_http(true);
173 if remote_store_options.is_empty() {
174 let http_store = object_store::http::HttpBuilder::new()
175 .with_url(url)
176 .with_client_options(client_options)
177 .with_retry(retry_config)
178 .build()?;
179 Ok(Box::new(http_store))
180 } else if Url::parse(&url)?.scheme() == "gs" {
181 let url = Url::parse(&url)?;
182 let mut builder = object_store::gcp::GoogleCloudStorageBuilder::new()
183 .with_url(url.as_str())
184 .with_retry(retry_config)
185 .with_client_options(client_options);
186 for (key, value) in remote_store_options {
187 builder = builder.with_config(GoogleConfigKey::from_str(&key)?, value);
188 }
189 Ok(Box::new(builder.build()?))
190 } else {
191 let url = Url::parse(&url)?;
192 let mut builder = object_store::aws::AmazonS3Builder::new()
193 .with_url(url.as_str())
194 .with_retry(retry_config)
195 .with_client_options(client_options);
196 for (key, value) in remote_store_options {
197 builder = builder.with_config(AmazonS3ConfigKey::from_str(&key)?, value);
198 }
199 Ok(Box::new(builder.build()?))
200 }
201}
202
203/// Creates a new [`ExponentialBackoff`] instance based on the configured
204/// template.
205///
206/// Returns a fresh backoff instance that has been reset to its initial
207/// state, ensuring consistent retry behavior for each new operation.
208pub(crate) fn reset_backoff(backoff: &ExponentialBackoff) -> ExponentialBackoff {
209 let mut backoff = backoff.clone();
210 backoff.reset();
211 backoff
212}