1#![allow(clippy::disallowed_types)] use std::collections::HashMap;
9#[cfg(any(feature = "aws", feature = "gcs", feature = "azure"))]
10use std::hash::Hash;
11use std::path::PathBuf;
12#[cfg(any(feature = "aws", feature = "gcs", feature = "azure"))]
13use std::str::FromStr;
14use std::sync::Arc;
15
16#[cfg(feature = "gcs")]
17use crate::gcs_credentials::{configure_gcs_workload_identity, gcs_workload_identity_provider};
18use crate::storage_auth::{classify_storage_auth_source, AuthSource};
19use crate::storage_location::{StorageConsumer, StorageLocation, StorageProvider};
20use object_store::ObjectStore;
21
22#[derive(Debug, thiserror::Error)]
24pub enum ObjectStoreBuilderError {
25 #[error("invalid object store URL: {0}")]
27 InvalidUrl(String),
28
29 #[error("object store build error: {0}")]
31 Build(String),
32}
33
34impl From<object_store::Error> for ObjectStoreBuilderError {
35 fn from(e: object_store::Error) -> Self {
36 Self::Build(e.to_string())
37 }
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
42pub enum CheckpointStorageScope {
43 Volatile,
45 NodeDurable,
47 ClusterShared,
49}
50
51impl CheckpointStorageScope {
52 #[must_use]
54 pub fn for_url(url: &str) -> Self {
55 let Ok(location) = StorageLocation::parse(url) else {
56 return Self::Volatile;
57 };
58 match location.provider {
59 StorageProvider::Local => Self::NodeDurable,
60 StorageProvider::AwsS3 | StorageProvider::AzureAdls | StorageProvider::Gcs => {
61 Self::ClusterShared
62 }
63 }
64 }
65
66 #[must_use]
68 pub const fn satisfies(self, required: Self) -> bool {
69 self as u8 >= required as u8
70 }
71}
72
73#[allow(clippy::implicit_hasher)]
93pub fn build_object_store(
94 url: &str,
95 options: &HashMap<String, String>,
96) -> Result<Arc<dyn ObjectStore>, ObjectStoreBuilderError> {
97 let mut location = StorageLocation::parse(url)
98 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
99 if location.provider == StorageProvider::Local {
100 return build_local_file_system(url);
101 }
102 let mut environment = provider_environment(location.provider);
103 if let Some(endpoint) = configured_endpoint(location.provider, options, &environment) {
104 location = location
105 .with_endpoint_override(endpoint)
106 .map_err(|error| ObjectStoreBuilderError::Build(error.to_string()))?;
107 }
108 let adapted = location
109 .adapt(StorageConsumer::ObjectStore)
110 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
111 let parsed = url::Url::parse(&adapted.url)
112 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
113 #[cfg(any(feature = "aws", feature = "gcs", feature = "azure"))]
114 validate_explicit_options(location.provider, options)?;
115 #[cfg(feature = "azure")]
116 validate_url_derived_azure_options(options, &adapted.derived_options)?;
117
118 environment.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
119 let auth_source = checkpoint_auth_source(location.provider, options);
120 #[cfg(feature = "gcs")]
121 let workload_identity = if location.provider == StorageProvider::Gcs {
122 gcs_workload_identity_provider(options, &|name| std::env::var(name).ok())
123 .map_err(|error| ObjectStoreBuilderError::Build(error.to_string()))?
124 } else {
125 None
126 };
127 #[cfg(feature = "gcs")]
128 let auth_source = if workload_identity.is_some() {
129 AuthSource::WorkloadIdentity
130 } else {
131 auth_source
132 };
133 let endpoint_transport = location
134 .endpoint_override
135 .as_ref()
136 .map_or("default", |endpoint| endpoint.scheme());
137 tracing::debug!(
138 operation_class = "client-build",
139 provider = location.provider.name(),
140 endpoint_class = location.endpoint_class().name(),
141 endpoint_transport,
142 allow_http = configured_allow_http(location.provider, options, &environment),
143 auth_source = %auth_source,
144 "building checkpoint object-store client"
145 );
146
147 let resolved = environment
150 .into_iter()
151 .chain(
152 options
153 .iter()
154 .map(|(key, value)| (key.clone(), value.clone())),
155 )
156 .chain(adapted.derived_options)
157 .collect::<Vec<_>>();
158 #[cfg(feature = "gcs")]
159 let result = if let Some(credentials) = workload_identity {
160 build_gcs_workload_identity_store(&parsed, resolved, credentials)
161 } else {
162 object_store::parse_url_opts(&parsed, resolved)
163 };
164 #[cfg(not(feature = "gcs"))]
165 let result = object_store::parse_url_opts(&parsed, resolved);
166 let (store, prefix) = result.map_err(|_| {
167 ObjectStoreBuilderError::Build(format!(
168 "failed to construct the {} client; verify the configured option values and downstream credential chain",
169 location.provider.name()
170 ))
171 })?;
172 let store: Arc<dyn ObjectStore> = Arc::from(store);
173 if prefix.as_ref().is_empty() {
174 Ok(store)
175 } else {
176 Ok(Arc::new(object_store::prefix::PrefixStore::new(
177 store, prefix,
178 )))
179 }
180}
181
182#[cfg(feature = "gcs")]
183fn build_gcs_workload_identity_store(
184 url: &url::Url,
185 options: Vec<(String, String)>,
186 credentials: object_store::gcp::GcpCredentialProvider,
187) -> object_store::Result<(Box<dyn ObjectStore>, object_store::path::Path)> {
188 let builder = object_store::gcp::GoogleCloudStorageBuilder::new().with_url(url.to_string());
189 let store = configure_gcs_workload_identity(builder, options, credentials).build()?;
190 let (_, prefix) = object_store::ObjectStoreScheme::parse(url)?;
191 let prefix = object_store::path::Path::parse(prefix)?;
192 Ok((Box::new(store), prefix))
193}
194
195pub fn durable_local_object_store(
201 root: impl AsRef<std::path::Path>,
202) -> object_store::Result<Arc<dyn ObjectStore>> {
203 Ok(Arc::new(
204 crate::durable_local_store::DurableLocalObjectStore::new(root)?,
205 ))
206}
207
208pub fn owned_durable_local_object_store(
216 root: impl AsRef<std::path::Path>,
217 owner: &Arc<std::fs::File>,
218) -> object_store::Result<Arc<dyn ObjectStore>> {
219 Ok(Arc::new(
220 crate::durable_local_store::DurableLocalObjectStore::new(root)?.with_namespace_owner(owner),
221 ))
222}
223
224fn provider_environment_key(provider: StorageProvider, key: &str) -> bool {
225 match provider {
226 StorageProvider::AwsS3 => key.starts_with("AWS_"),
227 StorageProvider::Gcs => key.starts_with("GOOGLE_") || key == "SERVICE_ACCOUNT",
228 StorageProvider::AzureAdls => key.starts_with("AZURE_") || key == "IDENTITY_ENDPOINT",
229 StorageProvider::Local => false,
230 }
231}
232
233fn provider_environment(provider: StorageProvider) -> Vec<(String, String)> {
234 std::env::vars_os()
235 .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?)))
236 .filter(|(key, _)| provider_environment_key(provider, key))
237 .collect()
238}
239
240fn checkpoint_auth_source(
241 provider: StorageProvider,
242 options: &HashMap<String, String>,
243) -> AuthSource {
244 let classified =
245 classify_storage_auth_source(provider, options, &|name| std::env::var(name).ok());
246 if provider == StorageProvider::AwsS3 && classified == AuthSource::Profile {
247 AuthSource::Unknown
250 } else {
251 classified
252 }
253}
254
255fn configured_endpoint<'a>(
256 provider: StorageProvider,
257 options: &'a HashMap<String, String>,
258 environment: &'a [(String, String)],
259) -> Option<&'a str> {
260 let (option_keys, environment_keys): (&[&str], &[&str]) = match provider {
261 StorageProvider::AwsS3 => (
262 &[
263 "aws_endpoint_url_s3",
264 "aws_endpoint",
265 "aws_endpoint_url",
266 "endpoint",
267 "endpoint_url",
268 ],
269 &["AWS_ENDPOINT_URL_S3", "AWS_ENDPOINT_URL", "AWS_ENDPOINT"],
270 ),
271 StorageProvider::AzureAdls => (
272 &["azure_storage_endpoint", "azure_endpoint", "endpoint"],
273 &["AZURE_STORAGE_ENDPOINT"],
274 ),
275 StorageProvider::Gcs => (
276 &["google_base_url", "base_url"],
277 &["GOOGLE_BASE_URL", "GOOGLE_ENDPOINT_URL"],
278 ),
279 StorageProvider::Local => (&[], &[]),
280 };
281 configured_option(options, option_keys)
282 .or_else(|| configured_environment(environment, environment_keys))
283}
284
285fn configured_allow_http(
286 provider: StorageProvider,
287 options: &HashMap<String, String>,
288 environment: &[(String, String)],
289) -> bool {
290 let (option_keys, environment_keys): (&[&str], &[&str]) = match provider {
291 StorageProvider::AwsS3 => (&["aws_allow_http", "allow_http"], &["AWS_ALLOW_HTTP"]),
292 StorageProvider::AzureAdls => (&["azure_allow_http", "allow_http"], &["AZURE_ALLOW_HTTP"]),
293 StorageProvider::Gcs => (&["google_allow_http", "allow_http"], &["GOOGLE_ALLOW_HTTP"]),
294 StorageProvider::Local => (&[], &[]),
295 };
296 configured_option(options, option_keys)
297 .or_else(|| configured_environment(environment, environment_keys))
298 .is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
299}
300
301fn configured_option<'a>(options: &'a HashMap<String, String>, keys: &[&str]) -> Option<&'a str> {
302 keys.iter().find_map(|key| {
303 options.iter().find_map(|(candidate, value)| {
304 (candidate.eq_ignore_ascii_case(key) && !value.trim().is_empty())
305 .then_some(value.as_str())
306 })
307 })
308}
309
310fn configured_environment<'a>(
311 environment: &'a [(String, String)],
312 keys: &[&str],
313) -> Option<&'a str> {
314 keys.iter().find_map(|key| {
315 environment.iter().find_map(|(candidate, value)| {
316 (candidate == key && !value.trim().is_empty()).then_some(value.as_str())
317 })
318 })
319}
320
321#[cfg(any(feature = "aws", feature = "gcs", feature = "azure"))]
322fn validate_explicit_options(
323 provider: StorageProvider,
324 options: &HashMap<String, String>,
325) -> Result<(), ObjectStoreBuilderError> {
326 match provider {
327 #[cfg(feature = "aws")]
328 StorageProvider::AwsS3 => {
329 validate_typed_options::<object_store::aws::AmazonS3ConfigKey>("S3", options)
330 }
331 #[cfg(feature = "gcs")]
332 StorageProvider::Gcs => {
333 validate_typed_options::<object_store::gcp::GoogleConfigKey>("GCS", options)
334 }
335 #[cfg(feature = "azure")]
336 StorageProvider::AzureAdls => {
337 validate_typed_options::<object_store::azure::AzureConfigKey>("Azure", options)
338 }
339 _ => Ok(()),
340 }
341}
342
343#[cfg(any(feature = "aws", feature = "gcs", feature = "azure"))]
344fn validate_typed_options<K>(
345 provider: &str,
346 options: &HashMap<String, String>,
347) -> Result<(), ObjectStoreBuilderError>
348where
349 K: FromStr + Eq + Hash,
350 K::Err: std::fmt::Display,
351{
352 let mut canonical = HashMap::<K, &str>::with_capacity(options.len());
353 for key in options.keys() {
354 let parsed = key.to_ascii_lowercase().parse::<K>().map_err(|error| {
355 ObjectStoreBuilderError::Build(format!(
356 "invalid {provider} storage option '{key}': {error}"
357 ))
358 })?;
359 if let Some(previous) = canonical.insert(parsed, key) {
360 return Err(ObjectStoreBuilderError::Build(format!(
361 "{provider} storage options '{previous}' and '{key}' configure the same setting"
362 )));
363 }
364 }
365 Ok(())
366}
367
368#[cfg(feature = "azure")]
369fn validate_url_derived_azure_options(
370 explicit: &HashMap<String, String>,
371 derived: &[(String, String)],
372) -> Result<(), ObjectStoreBuilderError> {
373 use object_store::azure::AzureConfigKey;
374
375 for (derived_key, derived_value) in derived {
376 let canonical = derived_key
377 .parse::<AzureConfigKey>()
378 .map_err(|error| ObjectStoreBuilderError::Build(error.to_string()))?;
379 for (explicit_key, explicit_value) in explicit {
380 let explicit_canonical = explicit_key
381 .to_ascii_lowercase()
382 .parse::<AzureConfigKey>()
383 .map_err(|error| ObjectStoreBuilderError::Build(error.to_string()))?;
384 if explicit_canonical == canonical && explicit_value != derived_value {
385 return Err(ObjectStoreBuilderError::Build(format!(
386 "Azure storage option '{explicit_key}' conflicts with the URL authority"
387 )));
388 }
389 }
390 }
391 Ok(())
392}
393
394pub fn file_url_path(url: &str) -> Result<PathBuf, ObjectStoreBuilderError> {
400 let location = StorageLocation::parse(url)
401 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
402 if location.provider != StorageProvider::Local {
403 return Err(ObjectStoreBuilderError::InvalidUrl(
404 "expected a file URL".to_string(),
405 ));
406 }
407 let adapted = location
408 .adapt(StorageConsumer::ObjectStore)
409 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
410 let parsed = url::Url::parse(&adapted.url)
411 .map_err(|error| ObjectStoreBuilderError::InvalidUrl(error.to_string()))?;
412 let path = parsed.to_file_path().map_err(|()| {
413 ObjectStoreBuilderError::InvalidUrl("file URL is not a local filesystem path".to_string())
414 })?;
415 if !path.is_absolute() {
416 return Err(ObjectStoreBuilderError::InvalidUrl(
417 "file URL path must be absolute".to_string(),
418 ));
419 }
420 Ok(path)
421}
422
423#[must_use]
429pub fn is_absolute_local_file_url(url: &str) -> bool {
430 StorageLocation::parse(url).is_ok_and(|location| location.provider == StorageProvider::Local)
431}
432
433fn build_local_file_system(url: &str) -> Result<Arc<dyn ObjectStore>, ObjectStoreBuilderError> {
435 let path = file_url_path(url)?;
436 durable_local_object_store(path).map_err(Into::into)
437}
438
439#[cfg(test)]
440mod tests;