From 77fc442a4d097a4b7a4a0ae1d83b0357556385a9 Mon Sep 17 00:00:00 2001 From: Ddupg Date: Fri, 7 Aug 2026 15:17:25 +0800 Subject: [PATCH] feat(io): add retry support for OpenDAL stores --- .../src/object_store/dynamic_opendal.rs | 86 ++++- rust/lance-io/src/object_store/providers.rs | 11 + .../src/object_store/providers/aws.rs | 3 + .../src/object_store/providers/azure.rs | 28 +- .../src/object_store/providers/gcp.rs | 4 + .../src/object_store/providers/goosefs.rs | 6 +- .../src/object_store/providers/huggingface.rs | 31 +- .../src/object_store/providers/opendal.rs | 309 ++++++++++++++++++ .../src/object_store/providers/oss.rs | 10 +- .../src/object_store/providers/tencent.rs | 6 +- .../src/object_store/providers/tos.rs | 10 +- 11 files changed, 471 insertions(+), 33 deletions(-) create mode 100644 rust/lance-io/src/object_store/providers/opendal.rs diff --git a/rust/lance-io/src/object_store/dynamic_opendal.rs b/rust/lance-io/src/object_store/dynamic_opendal.rs index 50b15d2180b..d5fc3477760 100644 --- a/rust/lance-io/src/object_store/dynamic_opendal.rs +++ b/rust/lance-io/src/object_store/dynamic_opendal.rs @@ -252,12 +252,29 @@ impl OSObjectStore for DynamicOpenDalStore { #[cfg(test)] mod tests { use std::collections::HashMap; - use std::sync::Arc; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + use async_trait::async_trait; use opendal::{Operator, services::Memory}; use super::*; + use crate::object_store::providers::opendal::finish_opendal_operator; use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{StorageOptions, StorageOptionsProvider}; + + fn build_memory_store(config: HashMap) -> Result { + let storage_options = StorageOptions(config); + let operator = Operator::new(Memory::default()).map_err(|e| { + lance_core::Error::invalid_input(format!("Failed to create memory operator: {e:?}")) + })?; + Ok(OpendalStore::new(finish_opendal_operator( + operator, + storage_options.client_max_retries(), + ))) + } #[tokio::test] async fn test_dynamic_store_caches_by_normalized_config() { @@ -272,14 +289,7 @@ mod tests { HashMap::new(), accessor, |options| Ok(options.clone()), - |_| { - let operator = Operator::new(Memory::default()).map_err(|e| { - lance_core::Error::invalid_input(format!( - "Failed to create memory operator: {e:?}" - )) - })?; - Ok(OpendalStore::new(operator)) - }, + build_memory_store, ); let first = store @@ -294,6 +304,55 @@ mod tests { assert!(Arc::ptr_eq(&first, &second)); } + #[derive(Debug)] + struct ChangingRetryConfigProvider { + fetch_count: AtomicUsize, + } + + #[async_trait] + impl StorageOptionsProvider for ChangingRetryConfigProvider { + async fn fetch_storage_options(&self) -> Result>> { + let fetch_count = self.fetch_count.fetch_add(1, Ordering::SeqCst); + Ok(Some(HashMap::from([ + ( + "client_max_retries".to_string(), + (fetch_count + 1).to_string(), + ), + ("expires_at_millis".to_string(), "0".to_string()), + ]))) + } + + fn provider_id(&self) -> String { + "ChangingRetryConfigProvider".to_string() + } + } + + #[tokio::test] + async fn test_dynamic_store_rebuilds_when_retry_config_changes() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + ChangingRetryConfigProvider { + fetch_count: AtomicUsize::new(0), + }, + ))); + let store = DynamicOpenDalStore::new( + "memory", + HashMap::new(), + accessor, + |options| Ok(options.clone()), + build_memory_store, + ); + + let first = store + .current_store() + .await + .expect("first store should build"); + let second = store + .current_store() + .await + .expect("changed retry config should rebuild store"); + + assert!(!Arc::ptr_eq(&first, &second)); + } #[test] fn test_merge_options_preserves_protected_base_keys() { let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( @@ -310,14 +369,7 @@ mod tests { ]), accessor, |options| Ok(options.clone()), - |_| { - let operator = Operator::new(Memory::default()).map_err(|e| { - lance_core::Error::invalid_input(format!( - "Failed to create memory operator: {e:?}" - )) - })?; - Ok(OpendalStore::new(operator)) - }, + build_memory_store, ) .with_protected_keys(["bucket", "root"]); diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index d8d184e07da..a2ef33e48be 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -30,6 +30,17 @@ pub mod goosefs; pub mod huggingface; pub mod local; pub mod memory; +#[cfg(any( + feature = "aws", + feature = "azure", + feature = "gcp", + feature = "goosefs", + feature = "huggingface", + feature = "oss", + feature = "tencent", + feature = "tos" +))] +pub(in crate::object_store) mod opendal; #[cfg(feature = "oss")] pub mod oss; pub mod shared_memory; diff --git a/rust/lance-io/src/object_store/providers/aws.rs b/rust/lance-io/src/object_store/providers/aws.rs index b23bc4cbb52..8a61dc3f0ff 100644 --- a/rust/lance-io/src/object_store/providers/aws.rs +++ b/rust/lance-io/src/object_store/providers/aws.rs @@ -40,6 +40,8 @@ use crate::object_store::{ use lance_core::error::{Error, Result}; use lance_core::utils::parse::str_is_truthy; +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct AwsStoreProvider; @@ -179,6 +181,7 @@ impl AwsStoreProvider { let operator = Operator::from_iter::(config_map) .map_err(|e| Error::invalid_input(format!("Failed to create S3 operator: {:?}", e)))?; + let operator = finish_opendal_operator(operator, storage_options.client_max_retries()); Ok(Arc::new(OpendalStore::new(operator)) as Arc) } diff --git a/rust/lance-io/src/object_store/providers/azure.rs b/rust/lance-io/src/object_store/providers/azure.rs index 2ad922fa241..9924272b004 100644 --- a/rust/lance-io/src/object_store/providers/azure.rs +++ b/rust/lance-io/src/object_store/providers/azure.rs @@ -28,6 +28,8 @@ use crate::object_store::{ use lance_core::error::{Error, Result}; use lance_core::utils::parse::str_is_truthy; +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct AzureBlobStoreProvider; @@ -93,6 +95,7 @@ impl AzureBlobStoreProvider { // Start with all storage options as the config map // OpenDAL will handle environment variables through its default credentials chain let mut config_map = Self::normalize_opendal_azure_options(&storage_options.0); + let max_retries = storage_options.client_max_retries(); match base_path.scheme() { "az" => { @@ -108,9 +111,14 @@ impl AzureBlobStoreProvider { config_map.insert("root".to_string(), format!("/{}", prefix)); } - Operator::from_iter::(config_map).map_err(|e| { - Error::invalid_input(format!("Failed to create Azure Blob operator: {:?}", e)) - }) + Operator::from_iter::(config_map) + .map_err(|e| { + Error::invalid_input(format!( + "Failed to create Azure Blob operator: {:?}", + e + )) + }) + .map(|operator| finish_opendal_operator(operator, max_retries)) } "abfss" => { let filesystem = base_path.username(); @@ -136,12 +144,14 @@ impl AzureBlobStoreProvider { config_map.insert("root".to_string(), format!("/{}", root_path)); } - Operator::from_iter::(config_map).map_err(|e| { - Error::invalid_input(format!( - "Failed to create Azure DFS (ADLS Gen2) operator: {:?}", - e - )) - }) + Operator::from_iter::(config_map) + .map_err(|e| { + Error::invalid_input(format!( + "Failed to create Azure DFS (ADLS Gen2) operator: {:?}", + e + )) + }) + .map(|operator| finish_opendal_operator(operator, max_retries)) } _ => Err(Error::invalid_input(format!( "Unsupported Azure scheme: {}", diff --git a/rust/lance-io/src/object_store/providers/gcp.rs b/rust/lance-io/src/object_store/providers/gcp.rs index d64462a013d..0bee0b0f05e 100644 --- a/rust/lance-io/src/object_store/providers/gcp.rs +++ b/rust/lance-io/src/object_store/providers/gcp.rs @@ -29,6 +29,9 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; use lance_core::utils::parse::str_is_truthy; + +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct GcsStoreProvider; @@ -228,6 +231,7 @@ impl GcsStoreProvider { let operator = Operator::from_iter::(config_map) .map_err(|e| Error::invalid_input(format!("Failed to create GCS operator: {:?}", e)))?; + let operator = finish_opendal_operator(operator, storage_options.client_max_retries()); Ok(Arc::new(OpendalStore::new(operator)) as Arc) } diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index 90aed8fa49d..30b812d8a90 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -14,6 +14,8 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; +use super::opendal::finish_opendal_operator; + /// Default GooseFS Master gRPC port. const DEFAULT_GOOSEFS_PORT: u16 = 9200; @@ -385,7 +387,8 @@ impl GooseFsStoreProvider { impl ObjectStoreProvider for GooseFsStoreProvider { async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); - let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); Self::validate_storage_option_keys(&storage_options)?; @@ -444,6 +447,7 @@ impl ObjectStoreProvider for GooseFsStoreProvider { let operator = Operator::from_iter::(config_map).map_err(|e| { Error::invalid_input(format!("Failed to create GooseFS operator: {:?}", e)) })?; + let operator = finish_opendal_operator(operator, storage_options.client_max_retries()); // Wrap as object_store::ObjectStore via OpendalStore bridge let opendal_store = Arc::new(OpendalStore::new(operator)); diff --git a/rust/lance-io/src/object_store/providers/huggingface.rs b/rust/lance-io/src/object_store/providers/huggingface.rs index dac76d130ef..aae30e63be4 100644 --- a/rust/lance-io/src/object_store/providers/huggingface.rs +++ b/rust/lance-io/src/object_store/providers/huggingface.rs @@ -18,6 +18,8 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; +use super::opendal::finish_opendal_operator; + /// Hugging Face object store provider backed by OpenDAL. #[derive(Default, Debug)] pub struct HuggingfaceStoreProvider; @@ -142,10 +144,20 @@ fn normalize_hf_config(options: &HashMap) -> Result) -> Result { + let storage_options = StorageOptions(config_map); + let max_retries = storage_options.client_max_retries(); + let config_map = &storage_options.0; let repo_type = config_map .get("repo_type") .ok_or_else(|| Error::invalid_input("Huggingface repo_type is required"))?; @@ -178,6 +190,7 @@ fn build_hf_store(config_map: HashMap) -> Result { let operator = Operator::new(builder).map_err(|e| { Error::invalid_input(format!("Failed to create Huggingface operator: {:?}", e)) })?; + let operator = finish_opendal_operator(operator, max_retries); Ok(OpendalStore::new(operator)) } @@ -190,7 +203,8 @@ impl ObjectStoreProvider for HuggingfaceStoreProvider { } = parse_hf_url(&base_path)?; let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); - let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); let download_retry_count = storage_options.download_retry_count(); let mut base_options = build_hf_base_options(&repo_type, &repo_id, &storage_options); @@ -414,6 +428,21 @@ mod tests { assert!(err.to_string().contains("Expected true or false")); } + #[test] + fn storage_option_preserves_client_max_retries() { + let config = normalize_hf_config(&build_hf_base_options( + "dataset", + "acme/repo", + &crate::object_store::StorageOptions(HashMap::from([( + "CLIENT_MAX_RETRIES".to_string(), + "5".to_string(), + )])), + )) + .unwrap(); + + assert_eq!(config.get("client_max_retries").unwrap(), "5"); + } + #[test] fn storage_option_download_mode_rejects_invalid_value() { let err = normalize_hf_config(&build_hf_base_options( diff --git a/rust/lance-io/src/object_store/providers/opendal.rs b/rust/lance-io/src/object_store/providers/opendal.rs new file mode 100644 index 00000000000..4a86715178d --- /dev/null +++ b/rust/lance-io/src/object_store/providers/opendal.rs @@ -0,0 +1,309 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::time::Duration; + +use opendal::Operator; +use opendal::layers::RetryLayer; + +pub(in crate::object_store) fn finish_opendal_operator( + operator: Operator, + max_retries: usize, +) -> Operator { + if max_retries == 0 { + return operator; + } + + let retry_layer = RetryLayer::new() + .with_max_times(max_retries) + .with_min_delay(Duration::from_millis(100)) + .with_max_delay(Duration::from_secs(15)) + .with_factor(2.0) + .with_jitter(); + + operator.layer(retry_layer) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use bytes::Bytes; + use opendal::raw::oio; + use opendal::raw::{ + OpCopy, OpCreateDir, OpList, OpPresign, OpRead, OpRename, OpStat, OpWrite, RpCreateDir, + RpPresign, RpRename, RpStat, Service, ServiceInfo, + }; + use opendal::{ + Buffer, Capability, Error, ErrorKind, Metadata, MetadataBuilder, OperationContext, Result, + }; + use rstest::rstest; + + use super::*; + + #[derive(Debug, Clone)] + struct RetryTestService { + stat_attempts: Arc, + stat_failures: usize, + is_stat_error_temporary: bool, + write_attempts: Arc, + write_failures: usize, + write_payloads: Arc>>, + close_attempts: Arc, + close_failures: usize, + } + + impl RetryTestService { + fn new(stat_failures: usize, is_stat_error_temporary: bool) -> Self { + Self { + stat_attempts: Arc::new(AtomicUsize::new(0)), + stat_failures, + is_stat_error_temporary, + write_attempts: Arc::new(AtomicUsize::new(0)), + write_failures: 0, + write_payloads: Arc::new(Mutex::new(Vec::new())), + close_attempts: Arc::new(AtomicUsize::new(0)), + close_failures: 0, + } + } + + fn with_write_failures(mut self, write_failures: usize) -> Self { + self.write_failures = write_failures; + self + } + + fn with_close_failures(mut self, close_failures: usize) -> Self { + self.close_failures = close_failures; + self + } + + fn stat_attempts(&self) -> usize { + self.stat_attempts.load(Ordering::Relaxed) + } + + fn write_attempts(&self) -> usize { + self.write_attempts.load(Ordering::Relaxed) + } + + fn write_payloads(&self) -> Vec { + self.write_payloads.lock().unwrap().clone() + } + + fn close_attempts(&self) -> usize { + self.close_attempts.load(Ordering::Relaxed) + } + } + + impl Service for RetryTestService { + type Reader = (); + type Writer = RetryTestWriter; + type Lister = (); + type Deleter = (); + type Copier = (); + type Composer = (); + + fn info(&self) -> ServiceInfo { + ServiceInfo::with_scheme("retry_test") + } + + fn capability(&self) -> Capability { + Capability { + stat: true, + write: true, + write_can_multi: true, + ..Default::default() + } + } + + async fn create_dir( + &self, + _: &OperationContext, + _: &str, + _: OpCreateDir, + ) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + async fn stat(&self, _: &OperationContext, _: &str, _: OpStat) -> Result { + let attempt = self.stat_attempts.fetch_add(1, Ordering::Relaxed) + 1; + if attempt <= self.stat_failures { + let error = Error::new(ErrorKind::Unexpected, "injected stat failure"); + return if self.is_stat_error_temporary { + Err(error.set_temporary()) + } else { + Err(error) + }; + } + + Ok(RpStat::new(MetadataBuilder::file(0).build())) + } + + fn read(&self, _: &OperationContext, _: &str, _: OpRead) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + fn write(&self, _: &OperationContext, _: &str, _: OpWrite) -> Result { + Ok(RetryTestWriter { + write_attempts: self.write_attempts.clone(), + write_failures: self.write_failures, + write_payloads: self.write_payloads.clone(), + close_attempts: self.close_attempts.clone(), + close_failures: self.close_failures, + content_length: 0, + }) + } + + fn delete(&self, _: &OperationContext) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + fn list(&self, _: &OperationContext, _: &str, _: OpList) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + fn copy(&self, _: &OperationContext, _: &str, _: &str, _: OpCopy) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + async fn rename( + &self, + _: &OperationContext, + _: &str, + _: &str, + _: OpRename, + ) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + + async fn presign(&self, _: &OperationContext, _: &str, _: OpPresign) -> Result { + Err(Error::new( + ErrorKind::Unsupported, + "operation is not supported", + )) + } + } + + fn build_test_operator(service: RetryTestService, max_retries: usize) -> Operator { + let operator = Operator::from_parts(OperationContext::default(), Arc::new(service)); + finish_opendal_operator(operator, max_retries) + } + + #[derive(Debug)] + struct RetryTestWriter { + write_attempts: Arc, + write_failures: usize, + write_payloads: Arc>>, + close_attempts: Arc, + close_failures: usize, + content_length: usize, + } + + impl oio::Write for RetryTestWriter { + async fn write(&mut self, buffer: Buffer) -> Result<()> { + let attempt = self.write_attempts.fetch_add(1, Ordering::Relaxed) + 1; + self.write_payloads.lock().unwrap().push(buffer.to_bytes()); + if attempt <= self.write_failures { + return Err( + Error::new(ErrorKind::Unexpected, "injected write failure").set_temporary() + ); + } + + self.content_length += buffer.len(); + Ok(()) + } + + async fn close(&mut self) -> Result { + let attempt = self.close_attempts.fetch_add(1, Ordering::Relaxed) + 1; + if attempt <= self.close_failures { + return Err( + Error::new(ErrorKind::Unexpected, "injected close failure").set_temporary() + ); + } + + Ok(MetadataBuilder::file(self.content_length as u64).build()) + } + + async fn abort(&mut self) -> Result<()> { + Ok(()) + } + } + + #[rstest] + #[case::recovers(2, true, 3, 3, None)] + #[case::persistent(1, false, 3, 1, Some(false))] + #[case::disabled(1, true, 0, 1, Some(true))] + #[case::exhausted(usize::MAX, true, 3, 4, Some(false))] + #[tokio::test(start_paused = true)] + async fn test_stat_retry( + #[case] failures: usize, + #[case] is_temporary: bool, + #[case] max_retries: usize, + #[case] expected_attempts: usize, + #[case] expected_temporary: Option, + ) { + let service = RetryTestService::new(failures, is_temporary); + let operator = build_test_operator(service.clone(), max_retries); + + let result = operator.stat("file").await; + assert_eq!(service.stat_attempts(), expected_attempts); + match expected_temporary { + None => { + result.expect("stat should succeed"); + } + Some(expected) => { + let error = result.expect_err("stat should fail"); + assert_eq!(error.is_temporary(), expected); + assert_eq!(error.is_persistent(), !expected); + } + } + } + + #[tokio::test(start_paused = true)] + async fn test_write_retry() { + let service = RetryTestService::new(0, true).with_write_failures(1); + let operator = build_test_operator(service.clone(), 3); + + operator + .write("file", "data") + .await + .expect("write should succeed"); + assert_eq!(service.write_attempts(), 2); + assert_eq!( + service.write_payloads(), + vec![Bytes::from_static(b"data"), Bytes::from_static(b"data")] + ); + assert_eq!(service.close_attempts(), 1); + } + + #[tokio::test(start_paused = true)] + async fn test_close_retry() { + let service = RetryTestService::new(0, true).with_close_failures(1); + let operator = build_test_operator(service.clone(), 3); + + operator + .write("file", "data") + .await + .expect("write should succeed"); + assert_eq!(service.close_attempts(), 2); + } +} diff --git a/rust/lance-io/src/object_store/providers/oss.rs b/rust/lance-io/src/object_store/providers/oss.rs index 0adf52db0a2..4ee7b784a15 100644 --- a/rust/lance-io/src/object_store/providers/oss.rs +++ b/rust/lance-io/src/object_store/providers/oss.rs @@ -16,6 +16,8 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct OssStoreProvider; @@ -94,8 +96,11 @@ impl OssStoreProvider { } fn build_oss_store(config_map: HashMap) -> Result { - let operator = Operator::from_iter::(config_map) + let storage_options = StorageOptions(config_map); + let max_retries = storage_options.client_max_retries(); + let operator = Operator::from_iter::(storage_options.0) .map_err(|e| Error::invalid_input(format!("Failed to create OSS operator: {:?}", e)))?; + let operator = finish_opendal_operator(operator, max_retries); Ok(OpendalStore::new(operator)) } @@ -105,7 +110,8 @@ impl OssStoreProvider { impl ObjectStoreProvider for OssStoreProvider { async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); - let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); let base_options = Self::base_oss_options(&base_path, &storage_options)?; let accessor = params.get_accessor(); diff --git a/rust/lance-io/src/object_store/providers/tencent.rs b/rust/lance-io/src/object_store/providers/tencent.rs index 9ad0a91765d..0654eeb094f 100644 --- a/rust/lance-io/src/object_store/providers/tencent.rs +++ b/rust/lance-io/src/object_store/providers/tencent.rs @@ -14,6 +14,8 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct TencentStoreProvider; @@ -21,7 +23,8 @@ pub struct TencentStoreProvider; impl ObjectStoreProvider for TencentStoreProvider { async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); - let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); let bucket = base_path .host_str() @@ -81,6 +84,7 @@ impl ObjectStoreProvider for TencentStoreProvider { let operator = Operator::from_iter::(config_map) .map_err(|e| Error::invalid_input(format!("Failed to create COS operator: {:?}", e)))?; + let operator = finish_opendal_operator(operator, storage_options.client_max_retries()); let opendal_store = Arc::new(OpendalStore::new(operator)); diff --git a/rust/lance-io/src/object_store/providers/tos.rs b/rust/lance-io/src/object_store/providers/tos.rs index 9f558e92550..2784a3e5c39 100644 --- a/rust/lance-io/src/object_store/providers/tos.rs +++ b/rust/lance-io/src/object_store/providers/tos.rs @@ -16,6 +16,8 @@ use crate::object_store::{ }; use lance_core::error::{Error, Result}; +use super::opendal::finish_opendal_operator; + #[derive(Default, Debug)] pub struct TosStoreProvider; @@ -100,8 +102,11 @@ impl TosStoreProvider { } fn build_tos_store(config_map: HashMap) -> Result { - let operator = Operator::from_iter::(config_map) + let storage_options = StorageOptions(config_map); + let max_retries = storage_options.client_max_retries(); + let operator = Operator::from_iter::(storage_options.0) .map_err(|e| Error::invalid_input(format!("Failed to create TOS operator: {:?}", e)))?; + let operator = finish_opendal_operator(operator, max_retries); Ok(OpendalStore::new(operator)) } @@ -111,7 +116,8 @@ impl TosStoreProvider { impl ObjectStoreProvider for TosStoreProvider { async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); - let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + let storage_options = + StorageOptions::new(params.storage_options().cloned().unwrap_or_default()); let base_options = Self::base_tos_options(&base_path, &storage_options)?; let accessor = params.get_accessor();