1use std::borrow::Cow;
7use std::collections::{HashMap, HashSet};
8use std::ops::Range;
9#[cfg(unix)]
10use std::os::unix::fs::PermissionsExt;
11use std::pin::Pin;
12use std::str::FromStr;
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15
16use ::tracing::{Span, field::Empty, instrument};
17use async_trait::async_trait;
18use bytes::Bytes;
19use chrono::{DateTime, Utc};
20use futures::{FutureExt, Stream};
21use futures::{StreamExt, TryStreamExt, future, stream::BoxStream};
22use lance_core::deepsize::DeepSizeOf;
23use lance_core::error::LanceOptionExt;
24use lance_core::utils::parse::{parse_env_as_bool, str_is_truthy};
25use list_retry::ListRetryStream;
26use object_store::DynObjectStore;
27use object_store::ObjectStoreExt as OSObjectStoreExt;
28#[cfg(feature = "aws")]
29use object_store::aws::AwsCredentialProvider;
30use object_store::list::PaginatedListStore;
31#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
32use object_store::{ClientOptions, HeaderMap, HeaderValue};
33use object_store::{
34 ListResult, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions, PutPayload,
35 path::Path,
36};
37use providers::local::FileStoreProvider;
38use providers::memory::MemoryStoreProvider;
39use tokio::io::AsyncWriteExt;
40use url::Url;
41
42use super::local::LocalObjectReader;
43#[cfg(target_os = "linux")]
44use crate::uring::{UringCurrentThreadReader, UringReader};
45#[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
46pub(crate) mod dynamic_credentials;
47#[cfg(any(feature = "oss", feature = "huggingface", feature = "tos"))]
48pub(crate) mod dynamic_opendal;
49mod list_retry;
50#[cfg(feature = "metrics")]
51pub mod metrics;
52#[cfg(any(
53 feature = "aws",
54 feature = "gcp",
55 feature = "azure",
56 feature = "oss",
57 feature = "tencent",
58 feature = "huggingface",
59 feature = "tos",
60 feature = "goosefs",
61))]
62pub(crate) mod opendal_store;
63pub mod providers;
64pub(crate) mod read_dir;
65pub mod storage_options;
66#[cfg(test)]
67pub(crate) mod test_utils;
68pub mod throttle;
69mod tracing;
70use crate::object_reader::SmallReader;
71use crate::object_writer::{LocalWriter, WriteResult};
72use crate::traits::{WriteExt, Writer};
73use crate::utils::tracking_store::{IOTracker, IoStats};
74use crate::{object_reader::CloudObjectReader, object_writer::ObjectWriter, traits::Reader};
75use lance_core::{Error, Result};
76
77pub const DEFAULT_LOCAL_IO_PARALLELISM: usize = 8;
82pub const DEFAULT_CLOUD_IO_PARALLELISM: usize = 64;
84
85const SERVER_SIDE_COPY_ENABLED_ENV: &str = "LANCE_IO_SERVER_SIDE_COPY_ENABLED";
86
87const DEFAULT_LOCAL_BLOCK_SIZE: usize = 4 * 1024; #[cfg(any(
89 feature = "aws",
90 feature = "gcp",
91 feature = "azure",
92 feature = "oss",
93 feature = "tencent",
94 feature = "huggingface",
95 feature = "tos",
96 feature = "goosefs",
97))]
98const DEFAULT_CLOUD_BLOCK_SIZE: usize = 64 * 1024; pub static DEFAULT_MAX_IOP_SIZE: std::sync::LazyLock<u64> = std::sync::LazyLock::new(|| {
101 std::env::var("LANCE_MAX_IOP_SIZE")
102 .map(|val| val.parse().unwrap())
103 .unwrap_or(16 * 1024 * 1024)
104});
105
106pub const DEFAULT_DOWNLOAD_RETRY_COUNT: usize = 3;
107
108#[derive(Debug)]
109struct StreamCopyError {
110 stage: &'static str,
111 source_path: String,
112 destination_path: String,
113 source: Box<dyn std::error::Error + Send + Sync>,
114}
115
116impl std::fmt::Display for StreamCopyError {
117 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118 write!(
119 f,
120 "multipart_stream_copy failed during {} from {} to {}: {}",
121 self.stage, self.source_path, self.destination_path, self.source
122 )
123 }
124}
125
126impl std::error::Error for StreamCopyError {
127 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
128 Some(self.source.as_ref())
129 }
130}
131
132fn stream_copy_error(
133 stage: &'static str,
134 source_path: &Path,
135 destination_path: &Path,
136 source: impl std::error::Error + Send + Sync + 'static,
137) -> Error {
138 Error::io_source(Box::new(StreamCopyError {
139 stage,
140 source_path: source_path.to_string(),
141 destination_path: destination_path.to_string(),
142 source: Box::new(source),
143 }))
144}
145
146pub use providers::{ObjectStoreProvider, ObjectStoreRegistry};
147pub use read_dir::ReadDirOptions;
148pub use storage_options::{
149 BASE_SCOPED_OPTION_PREFIX, BaseScopedStorageOptionsProvider, EXPIRES_AT_MILLIS_KEY,
150 LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, StorageOptionsAccessor,
151 StorageOptionsProvider, has_base_scoped_options, parse_base_scoped_key,
152 resolve_base_scoped_options,
153};
154
155#[async_trait]
156pub trait ObjectStoreExt {
157 async fn exists(&self, path: &Path) -> Result<bool>;
159
160 fn read_dir_all<'a, 'b>(
164 &'a self,
165 dir_path: impl Into<&'b Path> + Send,
166 unmodified_since: Option<DateTime<Utc>>,
167 ) -> BoxStream<'a, Result<ObjectMeta>>;
168}
169
170#[async_trait]
171pub(super) trait LocalDirOperations: std::fmt::Debug + Send + Sync {
172 async fn remove_dir_all(&self, path: &Path) -> Result<()>;
173}
174
175#[async_trait]
176impl<O: OSObjectStore + ?Sized> ObjectStoreExt for O {
177 fn read_dir_all<'a, 'b>(
178 &'a self,
179 dir_path: impl Into<&'b Path> + Send,
180 unmodified_since: Option<DateTime<Utc>>,
181 ) -> BoxStream<'a, Result<ObjectMeta>> {
182 let output = self.list(Some(dir_path.into())).map_err(|e| e.into());
183 if let Some(unmodified_since_val) = unmodified_since {
184 output
185 .try_filter(move |file| future::ready(file.last_modified <= unmodified_since_val))
186 .boxed()
187 } else {
188 output.boxed()
189 }
190 }
191
192 async fn exists(&self, path: &Path) -> Result<bool> {
193 match self.head(path).await {
194 Ok(_) => Ok(true),
195 Err(object_store::Error::NotFound { path: _, source: _ }) => Ok(false),
196 Err(e) => Err(e.into()),
197 }
198 }
199}
200
201#[derive(Clone)]
203pub struct ObjectStore {
204 pub inner: Arc<dyn OSObjectStore>,
206 local_dir_operations: Option<Arc<dyn LocalDirOperations>>,
208 scheme: String,
209 block_size: usize,
210 max_iop_size: u64,
211 pub use_constant_size_upload_parts: bool,
214 pub list_is_lexically_ordered: bool,
217 io_parallelism: usize,
218 download_retry_count: usize,
220 io_tracker: IOTracker,
222 pub store_prefix: String,
226 pub(crate) paginated_lister: Option<Arc<dyn PaginatedListStore>>,
229}
230
231impl std::fmt::Debug for ObjectStore {
233 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
234 f.debug_struct("ObjectStore")
235 .field("inner", &self.inner)
236 .field("scheme", &self.scheme)
237 .field("block_size", &self.block_size)
238 .field("max_iop_size", &self.max_iop_size)
239 .field(
240 "use_constant_size_upload_parts",
241 &self.use_constant_size_upload_parts,
242 )
243 .field("list_is_lexically_ordered", &self.list_is_lexically_ordered)
244 .field("io_parallelism", &self.io_parallelism)
245 .field("download_retry_count", &self.download_retry_count)
246 .field("io_tracker", &self.io_tracker)
247 .field("store_prefix", &self.store_prefix)
248 .field("paginated_lister", &self.paginated_lister.is_some())
249 .finish()
250 }
251}
252
253impl DeepSizeOf for ObjectStore {
254 fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
255 self.scheme.deep_size_of_children(context) + self.block_size.deep_size_of_children(context)
260 }
261}
262
263impl std::fmt::Display for ObjectStore {
264 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265 write!(f, "ObjectStore({})", self.scheme)
266 }
267}
268
269pub trait WrappingObjectStore: std::fmt::Debug + Send + Sync {
270 fn wrap(&self, store_prefix: &str, original: Arc<dyn OSObjectStore>) -> Arc<dyn OSObjectStore>;
275
276 fn wrap_paginated(
299 &self,
300 store_prefix: &str,
301 original: Arc<dyn PaginatedListStore>,
302 ) -> Option<Arc<dyn PaginatedListStore>>;
303}
304
305#[derive(Debug, Clone)]
306pub struct ChainedWrappingObjectStore {
307 wrappers: Vec<Arc<dyn WrappingObjectStore>>,
308}
309
310impl ChainedWrappingObjectStore {
311 pub fn new(wrappers: Vec<Arc<dyn WrappingObjectStore>>) -> Self {
312 Self { wrappers }
313 }
314
315 pub fn add_wrapper(&mut self, wrapper: Arc<dyn WrappingObjectStore>) {
316 self.wrappers.push(wrapper);
317 }
318}
319
320impl WrappingObjectStore for ChainedWrappingObjectStore {
321 fn wrap(&self, store_prefix: &str, original: Arc<dyn OSObjectStore>) -> Arc<dyn OSObjectStore> {
322 self.wrappers
323 .iter()
324 .fold(original, |acc, wrapper| wrapper.wrap(store_prefix, acc))
325 }
326
327 fn wrap_paginated(
330 &self,
331 store_prefix: &str,
332 original: Arc<dyn PaginatedListStore>,
333 ) -> Option<Arc<dyn PaginatedListStore>> {
334 self.wrappers.iter().try_fold(original, |acc, wrapper| {
335 wrapper.wrap_paginated(store_prefix, acc)
336 })
337 }
338}
339
340#[derive(Debug, Clone)]
343pub struct ObjectStoreParams {
344 pub block_size: Option<usize>,
345 #[deprecated(note = "Implement an ObjectStoreProvider instead")]
346 pub object_store: Option<(Arc<DynObjectStore>, Url)>,
347 pub s3_credentials_refresh_offset: Duration,
350 #[cfg(feature = "aws")]
351 pub aws_credentials: Option<AwsCredentialProvider>,
352 pub object_store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
353 pub storage_options_accessor: Option<Arc<StorageOptionsAccessor>>,
359 pub use_constant_size_upload_parts: bool,
364 pub list_is_lexically_ordered: Option<bool>,
365}
366
367impl Default for ObjectStoreParams {
368 fn default() -> Self {
369 #[allow(deprecated)]
370 Self {
371 object_store: None,
372 block_size: None,
373 s3_credentials_refresh_offset: Duration::from_secs(60),
374 #[cfg(feature = "aws")]
375 aws_credentials: None,
376 object_store_wrapper: None,
377 storage_options_accessor: None,
378 use_constant_size_upload_parts: false,
379 list_is_lexically_ordered: None,
380 }
381 }
382}
383
384impl ObjectStoreParams {
385 pub fn get_accessor(&self) -> Option<Arc<StorageOptionsAccessor>> {
387 self.storage_options_accessor.clone()
388 }
389
390 pub fn storage_options(&self) -> Option<&HashMap<String, String>> {
394 self.storage_options_accessor
395 .as_ref()
396 .and_then(|a| a.initial_storage_options())
397 }
398
399 pub fn resolved_block_size(&self) -> Result<Option<usize>> {
402 if self.block_size.is_some() {
403 return Ok(self.block_size);
404 }
405 match self.storage_options() {
406 Some(options) => StorageOptions(options.clone()).block_size(),
407 None => Ok(None),
408 }
409 }
410
411 pub fn scoped_to_base(&self, base_id: Option<u32>) -> Cow<'_, Self> {
418 let Some(accessor) = &self.storage_options_accessor else {
419 return Cow::Borrowed(self);
420 };
421 let scoped = accessor.scoped_to_base(base_id);
422 if Arc::ptr_eq(&scoped, accessor) {
423 Cow::Borrowed(self)
424 } else {
425 Cow::Owned(Self {
426 storage_options_accessor: Some(scoped),
427 ..self.clone()
428 })
429 }
430 }
431}
432
433fn wrapper_allocation_ptr(wrapper: &Arc<dyn WrappingObjectStore>) -> *const () {
434 Arc::as_ptr(wrapper) as *const ()
437}
438
439impl std::hash::Hash for ObjectStoreParams {
441 #[allow(deprecated)]
442 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
443 self.block_size.hash(state);
445 if let Some((store, url)) = &self.object_store {
446 Arc::as_ptr(store).hash(state);
447 url.hash(state);
448 }
449 self.s3_credentials_refresh_offset.hash(state);
450 #[cfg(feature = "aws")]
451 if let Some(aws_credentials) = &self.aws_credentials {
452 Arc::as_ptr(aws_credentials).hash(state);
453 }
454 if let Some(wrapper) = &self.object_store_wrapper {
455 wrapper_allocation_ptr(wrapper).hash(state);
456 }
457 if let Some(accessor) = &self.storage_options_accessor {
458 accessor.accessor_id().hash(state);
459 }
460 self.use_constant_size_upload_parts.hash(state);
461 self.list_is_lexically_ordered.hash(state);
462 }
463}
464
465impl Eq for ObjectStoreParams {}
467impl PartialEq for ObjectStoreParams {
468 #[allow(deprecated)]
469 fn eq(&self, other: &Self) -> bool {
470 #[cfg(feature = "aws")]
471 if self.aws_credentials.is_some() != other.aws_credentials.is_some() {
472 return false;
473 }
474
475 self.block_size == other.block_size
478 && self
479 .object_store
480 .as_ref()
481 .map(|(store, url)| (Arc::as_ptr(store), url))
482 == other
483 .object_store
484 .as_ref()
485 .map(|(store, url)| (Arc::as_ptr(store), url))
486 && self.s3_credentials_refresh_offset == other.s3_credentials_refresh_offset
487 && self
488 .object_store_wrapper
489 .as_ref()
490 .map(wrapper_allocation_ptr)
491 == other
492 .object_store_wrapper
493 .as_ref()
494 .map(wrapper_allocation_ptr)
495 && self
496 .storage_options_accessor
497 .as_ref()
498 .map(|a| a.accessor_id())
499 == other
500 .storage_options_accessor
501 .as_ref()
502 .map(|a| a.accessor_id())
503 && self.use_constant_size_upload_parts == other.use_constant_size_upload_parts
504 && self.list_is_lexically_ordered == other.list_is_lexically_ordered
505 }
506}
507
508pub fn uri_to_url(uri: &str) -> Result<Url> {
529 match Url::parse(uri) {
530 Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
531 local_path_to_url(uri)
533 }
534 Ok(url) => Ok(url),
535 Err(_) => local_path_to_url(uri),
536 }
537}
538
539fn expand_path(str_path: impl AsRef<str>) -> Result<std::path::PathBuf> {
540 let str_path = str_path.as_ref();
541 let expanded = expand_tilde_path(str_path).unwrap_or_else(|| str_path.into());
542
543 let mut expanded_path = path_abs::PathAbs::new(expanded)
544 .unwrap()
545 .as_path()
546 .to_path_buf();
547 if let Some(s) = expanded_path.as_path().to_str()
549 && s.is_empty()
550 {
551 expanded_path = std::env::current_dir()?;
552 }
553
554 Ok(expanded_path)
555}
556
557fn expand_tilde_path(path: &str) -> Option<std::path::PathBuf> {
558 let home_dir = std::env::home_dir()?;
559 if path == "~" {
560 return Some(home_dir);
561 }
562 if let Some(stripped) = path.strip_prefix("~/") {
563 return Some(home_dir.join(stripped));
564 }
565 #[cfg(windows)]
566 if let Some(stripped) = path.strip_prefix("~\\") {
567 return Some(home_dir.join(stripped));
568 }
569
570 None
571}
572
573fn local_path_to_url(str_path: &str) -> Result<Url> {
574 let expanded_path = expand_path(str_path)?;
575
576 Url::from_directory_path(expanded_path).map_err(|_| {
577 Error::invalid_input_source(format!("Invalid table location: '{}'", str_path).into())
578 })
579}
580
581#[cfg(feature = "huggingface")]
582fn parse_hf_repo_id(url: &Url) -> Result<String> {
583 let mut segments: Vec<String> = Vec::new();
585 if let Some(host) = url.host_str() {
586 segments.push(host.to_string());
587 }
588 segments.extend(
589 url.path()
590 .trim_start_matches('/')
591 .split('/')
592 .map(|s| s.to_string()),
593 );
594
595 if segments.len() < 2 {
596 return Err(Error::invalid_input(
597 "Huggingface URL must contain at least owner and repo",
598 ));
599 }
600
601 let repo_type_candidates = ["models", "datasets", "spaces"];
602 let (owner, repo_with_rev) = if repo_type_candidates.contains(&segments[0].as_str()) {
603 if segments.len() < 3 {
604 return Err(Error::invalid_input(
605 "Huggingface URL missing owner/repo after repo type",
606 ));
607 }
608 (segments[1].as_str(), segments[2].as_str())
609 } else {
610 (segments[0].as_str(), segments[1].as_str())
611 };
612
613 let repo = repo_with_rev
614 .split_once('@')
615 .map(|(r, _)| r)
616 .unwrap_or(repo_with_rev);
617 Ok(format!("{owner}/{repo}"))
618}
619
620impl ObjectStore {
621 pub async fn from_uri(uri: &str) -> Result<(Arc<Self>, Path)> {
629 let registry = Arc::new(ObjectStoreRegistry::default());
630
631 Self::from_uri_and_params(registry, uri, &ObjectStoreParams::default()).await
632 }
633
634 pub async fn from_uri_and_params(
638 registry: Arc<ObjectStoreRegistry>,
639 uri: &str,
640 params: &ObjectStoreParams,
641 ) -> Result<(Arc<Self>, Path)> {
642 Self::from_uri_and_params_impl(registry, uri, params, true).await
643 }
644
645 #[doc(hidden)]
650 pub async fn from_uri_and_params_uncached(
651 registry: Arc<ObjectStoreRegistry>,
652 uri: &str,
653 params: &ObjectStoreParams,
654 ) -> Result<(Arc<Self>, Path)> {
655 Self::from_uri_and_params_impl(registry, uri, params, false).await
656 }
657
658 async fn from_uri_and_params_impl(
659 registry: Arc<ObjectStoreRegistry>,
660 uri: &str,
661 params: &ObjectStoreParams,
662 use_registry_cache: bool,
663 ) -> Result<(Arc<Self>, Path)> {
664 #[allow(deprecated)]
665 if let Some((store, path)) = params.object_store.as_ref() {
666 let mut inner = store.clone();
667 let store_prefix =
668 registry.calculate_object_store_prefix(uri, params.storage_options())?;
669
670 let mut io_tracker = IOTracker::default();
671 meter_store(
672 &mut inner,
673 &mut io_tracker,
674 &metrics_base(&store_prefix, path),
675 );
676
677 if let Some(wrapper) = params.object_store_wrapper.as_ref() {
678 inner = wrapper.wrap(&store_prefix, inner);
679 }
680
681 let tracked_store = io_tracker.wrap("", inner);
683
684 let store = Self {
685 inner: tracked_store,
686 local_dir_operations: None,
687 scheme: path.scheme().to_string(),
688 block_size: params.resolved_block_size()?.unwrap_or(64 * 1024),
689 max_iop_size: *DEFAULT_MAX_IOP_SIZE,
690 use_constant_size_upload_parts: params.use_constant_size_upload_parts,
691 list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or_default(),
692 io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
693 download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT,
694 io_tracker,
695 store_prefix,
696 paginated_lister: None,
698 };
699 let path = Path::parse(path.path())?;
700 return Ok((Arc::new(store), path));
701 }
702 let url = uri_to_url(uri)?;
703
704 let store = if use_registry_cache {
705 registry.get_store(url.clone(), params).await?
706 } else {
707 registry.new_store(url.clone(), params).await?
708 };
709 let provider = registry.get_provider(url.scheme()).expect_ok()?;
711 let path = provider.extract_path(&url)?;
712
713 Ok((store, path))
714 }
715
716 pub fn extract_path_from_uri(registry: Arc<ObjectStoreRegistry>, uri: &str) -> Result<Path> {
730 let url = uri_to_url(uri)?;
731 let provider = registry
732 .get_provider(url.scheme())
733 .ok_or_else(|| Error::invalid_input(format!("Unknown scheme: {}", url.scheme())))?;
734 provider.extract_path(&url)
735 }
736
737 #[deprecated(note = "Use `from_uri` instead")]
738 pub fn from_path(str_path: &str) -> Result<(Arc<Self>, Path)> {
739 Self::from_uri_and_params(
740 Arc::new(ObjectStoreRegistry::default()),
741 str_path,
742 &Default::default(),
743 )
744 .now_or_never()
745 .unwrap()
746 }
747
748 pub fn local() -> Self {
750 let provider = FileStoreProvider;
751 provider
752 .new_store(Url::parse("file:///").unwrap(), &Default::default())
753 .now_or_never()
754 .unwrap()
755 .unwrap()
756 }
757
758 pub fn memory() -> Self {
760 let provider = MemoryStoreProvider;
761 provider
762 .new_store(Url::parse("memory:///").unwrap(), &Default::default())
763 .now_or_never()
764 .unwrap()
765 .unwrap()
766 }
767
768 pub fn is_local(&self) -> bool {
770 self.scheme == "file" || self.scheme == "file+uring"
771 }
772
773 pub fn has_direct_local_paths(&self) -> bool {
778 self.is_local() && self.store_prefix == self.scheme
779 }
780
781 pub fn is_cloud(&self) -> bool {
782 if self.is_local() || self.scheme == "memory" || self.scheme == "shared-memory" {
783 return false;
784 }
785 true
786 }
787
788 pub fn prefers_lite_scheduler(&self) -> bool {
793 self.scheme == "file+uring"
794 }
795
796 pub fn scheme(&self) -> &str {
797 &self.scheme
798 }
799
800 pub fn block_size(&self) -> usize {
801 self.block_size
802 }
803
804 pub fn max_iop_size(&self) -> u64 {
805 self.max_iop_size
806 }
807
808 pub fn io_parallelism(&self) -> usize {
815 std::env::var("LANCE_IO_THREADS")
816 .map(|val| val.parse::<usize>().unwrap())
817 .unwrap_or(self.io_parallelism)
818 .max(1)
819 }
820
821 pub fn io_tracker(&self) -> &IOTracker {
826 &self.io_tracker
827 }
828
829 pub fn io_stats_snapshot(&self) -> IoStats {
834 self.io_tracker.stats()
835 }
836
837 pub fn io_stats_incremental(&self) -> IoStats {
843 self.io_tracker.incremental_stats()
844 }
845
846 pub fn apply_wrapper(&mut self, wrapper: &dyn WrappingObjectStore) {
852 self.inner = wrapper.wrap(&self.store_prefix, self.inner.clone());
853 self.paginated_lister = self
854 .paginated_lister
855 .take()
856 .and_then(|lister| wrapper.wrap_paginated(&self.store_prefix, lister));
857 }
858
859 pub async fn open(&self, path: &Path) -> Result<Box<dyn Reader>> {
864 match self.scheme.as_str() {
865 "file" if self.has_direct_local_paths() => {
866 LocalObjectReader::open_with_tracker(
867 path,
868 self.block_size,
869 None,
870 Arc::new(self.io_tracker.clone()),
871 )
872 .await
873 }
874 #[cfg(target_os = "linux")]
875 "file+uring" => {
876 let use_current_thread = std::env::var("LANCE_URING_CURRENT_THREAD")
878 .map(|v| str_is_truthy(&v))
879 .unwrap_or(false);
880
881 if use_current_thread {
882 UringCurrentThreadReader::open(
883 path,
884 self.block_size,
885 None,
886 Arc::new(self.io_tracker.clone()),
887 )
888 .await
889 } else {
890 UringReader::open(
891 path,
892 self.block_size,
893 None,
894 Arc::new(self.io_tracker.clone()),
895 )
896 .await
897 }
898 }
899 _ => Ok(Box::new(
900 CloudObjectReader::new(
901 self.inner.clone(),
902 path.clone(),
903 self.block_size,
904 None,
905 self.download_retry_count,
906 )?
907 .with_io_parallelism(self.io_parallelism()),
908 )),
909 }
910 }
911
912 pub async fn open_with_size(&self, path: &Path, known_size: usize) -> Result<Box<dyn Reader>> {
918 if known_size <= self.block_size {
921 return Ok(Box::new(SmallReader::new(
922 self.inner.clone(),
923 path.clone(),
924 self.download_retry_count,
925 known_size,
926 )));
927 }
928
929 match self.scheme.as_str() {
930 "file" if self.has_direct_local_paths() => {
931 LocalObjectReader::open_with_tracker(
932 path,
933 self.block_size,
934 Some(known_size),
935 Arc::new(self.io_tracker.clone()),
936 )
937 .await
938 }
939 #[cfg(target_os = "linux")]
940 "file+uring" => {
941 let use_current_thread = std::env::var("LANCE_URING_CURRENT_THREAD")
943 .map(|v| str_is_truthy(&v))
944 .unwrap_or(false);
945
946 if use_current_thread {
947 UringCurrentThreadReader::open(
948 path,
949 self.block_size,
950 Some(known_size),
951 Arc::new(self.io_tracker.clone()),
952 )
953 .await
954 } else {
955 UringReader::open(
956 path,
957 self.block_size,
958 Some(known_size),
959 Arc::new(self.io_tracker.clone()),
960 )
961 .await
962 }
963 }
964 _ => Ok(Box::new(
965 CloudObjectReader::new(
966 self.inner.clone(),
967 path.clone(),
968 self.block_size,
969 Some(known_size),
970 self.download_retry_count,
971 )?
972 .with_io_parallelism(self.io_parallelism()),
973 )),
974 }
975 }
976
977 pub async fn create_local_writer(path: &std::path::Path) -> Result<ObjectWriter> {
979 let object_store = Self::local();
980 let absolute_path = expand_path(path.to_string_lossy())?;
981 let os_path = Path::from_absolute_path(absolute_path)?;
982 ObjectWriter::new(&object_store, &os_path).await
983 }
984
985 pub async fn open_local(path: &std::path::Path) -> Result<Box<dyn Reader>> {
987 let object_store = Self::local();
988 let absolute_path = expand_path(path.to_string_lossy())?;
989 let os_path = Path::from_absolute_path(absolute_path)?;
990 object_store.open(&os_path).await
991 }
992
993 pub async fn create(&self, path: &Path) -> Result<Box<dyn Writer>> {
995 match self.scheme.as_str() {
996 "file" if self.has_direct_local_paths() => {
997 let local_path = super::local::to_local_path(path);
998 let local_path = std::path::PathBuf::from(&local_path);
999 if let Some(parent) = local_path.parent() {
1000 tokio::fs::create_dir_all(parent).await?;
1001 }
1002 let parent = local_path
1003 .parent()
1004 .expect("file path must have parent")
1005 .to_owned();
1006 let named_temp = tokio::task::spawn_blocking(move || {
1007 #[cfg(unix)]
1008 {
1009 tempfile::Builder::new()
1011 .permissions(std::fs::Permissions::from_mode(0o666))
1012 .tempfile_in(parent)
1013 }
1014 #[cfg(not(unix))]
1015 tempfile::NamedTempFile::new_in(parent)
1016 })
1017 .await
1018 .map_err(|e| Error::io(format!("spawn_blocking failed: {}", e)))??;
1019 let (std_file, temp_path) = named_temp.into_parts();
1020 let file = tokio::fs::File::from_std(std_file);
1021 Ok(Box::new(LocalWriter::new(
1022 file,
1023 path.clone(),
1024 temp_path,
1025 Arc::new(self.io_tracker.clone()),
1026 )))
1027 }
1028 _ => Ok(Box::new(ObjectWriter::new(self, path).await?)),
1029 }
1030 }
1031
1032 pub async fn put(&self, path: &Path, content: &[u8]) -> Result<WriteResult> {
1034 let mut writer = self.create(path).await?;
1035 writer.write_all(content).await?;
1036 Writer::shutdown(writer.as_mut()).await
1037 }
1038
1039 pub async fn put_if_absent(
1048 &self,
1049 path: &Path,
1050 content: PutPayload,
1051 ) -> object_store::Result<()> {
1052 if self.scheme == "cos" {
1053 return Err(object_store::Error::NotSupported {
1054 source: "Tencent COS does not reliably enforce put-if-absent after bucket \
1055 versioning has ever been enabled"
1056 .into(),
1057 });
1058 }
1059
1060 if self.is_local() {
1061 let staging_path =
1062 Path::from(format!("{}.tmp.{}", path, uuid::Uuid::new_v4().simple()));
1063 self.inner.put(&staging_path, content).await?;
1064 let result = self.inner.rename_if_not_exists(&staging_path, path).await;
1065 if result.is_err()
1066 && let Err(error) = self.inner.delete(&staging_path).await
1067 {
1068 log::warn!(
1069 "Failed to remove staging object {} after atomic create failed: {}",
1070 staging_path,
1071 error
1072 );
1073 }
1074 result
1075 } else {
1076 self.inner
1077 .put_opts(
1078 path,
1079 content,
1080 PutOptions {
1081 mode: PutMode::Create,
1082 ..Default::default()
1083 },
1084 )
1085 .await
1086 .map(|_| ())
1087 }
1088 }
1089
1090 pub async fn delete(&self, path: &Path) -> Result<()> {
1091 self.inner.delete(path).await?;
1092 Ok(())
1093 }
1094
1095 const MAX_SINGLE_COPY_BYTES: u64 = 5 * 1024 * 1024 * 1024; pub async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
1100 let multipart_copy_fallback = matches!(self.scheme.as_str(), "s3" | "s3+ddb" | "gs");
1106 self.copy_impl(
1107 from,
1108 to,
1109 multipart_copy_fallback,
1110 Self::MAX_SINGLE_COPY_BYTES,
1111 )
1112 .await
1113 }
1114
1115 pub async fn copy_bulk(
1139 &self,
1140 source_path: &Path,
1141 destination_store: &Self,
1142 destination_path: &Path,
1143 ) -> Result<WriteResult> {
1144 self.copy_bulk_with_server_side_copy(
1145 source_path,
1146 destination_store,
1147 destination_path,
1148 self.uses_server_side_copy(destination_store),
1149 )
1150 .await
1151 }
1152
1153 fn uses_server_side_copy(&self, destination_store: &Self) -> bool {
1154 parse_env_as_bool(SERVER_SIDE_COPY_ENABLED_ENV, false)
1155 && self.can_server_side_copy_to(destination_store)
1156 }
1157
1158 async fn copy_bulk_with_server_side_copy(
1159 &self,
1160 source_path: &Path,
1161 destination_store: &Self,
1162 destination_path: &Path,
1163 server_side_copy_enabled: bool,
1164 ) -> Result<WriteResult> {
1165 if !server_side_copy_enabled || !self.can_server_side_copy_to(destination_store) {
1166 return self
1167 .copy_via_stream(source_path, destination_store, destination_path)
1168 .await;
1169 }
1170
1171 let source_size = self.size(source_path).await?;
1172 let result_size = usize::try_from(source_size).map_err(|source| {
1173 Error::io(format!(
1174 "server-side copy source size conversion failed from {source_path} to \
1175 {destination_path}: source_size={source_size}, error={source}"
1176 ))
1177 })?;
1178 destination_store
1179 .copy(source_path, destination_path)
1180 .await?;
1181 let destination_size = destination_store.size(destination_path).await?;
1182 if destination_size != source_size {
1183 return Err(Error::io(format!(
1184 "server-side copy destination size mismatch from {source_path} to \
1185 {destination_path}: source_size={source_size}, \
1186 destination_size={destination_size}"
1187 )));
1188 }
1189
1190 Ok(WriteResult {
1191 size: result_size,
1192 e_tag: None,
1193 })
1194 }
1195
1196 fn can_server_side_copy_to(&self, destination_store: &Self) -> bool {
1197 self.is_cloud()
1200 && destination_store.is_cloud()
1201 && Arc::ptr_eq(&self.inner, &destination_store.inner)
1202 }
1203
1204 #[instrument(
1227 name = "multipart_stream_copy",
1228 level = "info",
1229 skip(self, source_path, destination_store, destination_path),
1230 fields(
1231 source = %source_path,
1232 destination = %destination_path,
1233 source_size = Empty,
1234 read_chunk_size = Empty,
1235 multipart_part_size = crate::object_writer::initial_upload_size(),
1236 multipart_concurrency = crate::object_writer::max_upload_parallelism(),
1237 part_count = Empty,
1238 bytes_transferred = Empty,
1239 destination_size = Empty,
1240 validation = Empty,
1241 elapsed_ms = Empty,
1242 ),
1243 err
1244 )]
1245 pub async fn copy_via_stream(
1246 &self,
1247 source_path: &Path,
1248 destination_store: &Self,
1249 destination_path: &Path,
1250 ) -> Result<WriteResult> {
1251 let started_at = Instant::now();
1252 if self.has_direct_local_paths() && destination_store.has_direct_local_paths() {
1253 let source_size = std::fs::metadata(super::local::to_local_path(source_path))
1254 .map_err(|source| {
1255 let source = if source.kind() == std::io::ErrorKind::NotFound {
1256 Error::not_found(source_path.to_string())
1257 } else {
1258 Error::from(source)
1259 };
1260 stream_copy_error("source metadata", source_path, destination_path, source)
1261 })?
1262 .len();
1263 let source_size = usize::try_from(source_size).map_err(|source| {
1264 stream_copy_error(
1265 "source size conversion",
1266 source_path,
1267 destination_path,
1268 source,
1269 )
1270 })?;
1271 Span::current().record("source_size", source_size as u64);
1272
1273 let metrics = destination_store.io_tracker.begin_io("copy");
1274 let result = super::local::copy_file(source_path, destination_path);
1275 metrics.record(&result, source_size as u64);
1276 result.map_err(|source| {
1277 stream_copy_error(
1278 "local filesystem copy",
1279 source_path,
1280 destination_path,
1281 source,
1282 )
1283 })?;
1284
1285 let destination_size =
1286 destination_store
1287 .size(destination_path)
1288 .await
1289 .map_err(|source| {
1290 stream_copy_error(
1291 "destination validation",
1292 source_path,
1293 destination_path,
1294 source,
1295 )
1296 })?;
1297 Span::current().record("bytes_transferred", source_size as u64);
1298 Span::current().record("destination_size", destination_size);
1299 if destination_size != source_size as u64 {
1300 Span::current().record("validation", "failed");
1301 return Err(Error::io(format!(
1302 "multipart_stream_copy destination size mismatch from {source_path} to \
1303 {destination_path}: source_size={source_size}, \
1304 destination_size={destination_size}"
1305 )));
1306 }
1307
1308 Span::current().record("validation", "passed");
1309 Span::current().record("elapsed_ms", started_at.elapsed().as_millis() as u64);
1310 return Ok(WriteResult {
1311 size: source_size,
1312 e_tag: None,
1313 });
1314 }
1315
1316 let reader = self.open(source_path).await.map_err(|source| {
1317 stream_copy_error("source open", source_path, destination_path, source)
1318 })?;
1319 let source_size = reader.size().await.map_err(|source| {
1320 stream_copy_error("source metadata", source_path, destination_path, source)
1321 })?;
1322 Span::current().record("source_size", source_size as u64);
1323
1324 let mut writer = destination_store
1325 .create(destination_path)
1326 .await
1327 .map_err(|source| {
1328 stream_copy_error(
1329 "destination writer creation",
1330 source_path,
1331 destination_path,
1332 source,
1333 )
1334 })?;
1335 let read_chunk_size = usize::try_from(self.max_iop_size())
1336 .unwrap_or(usize::MAX)
1337 .max(1);
1338 Span::current().record("read_chunk_size", read_chunk_size as u64);
1339 let mut bytes_transferred = 0usize;
1340 if source_size > 0 {
1341 let first_range = 0..read_chunk_size.min(source_size);
1342 let mut current_range = first_range.clone();
1343 let mut current_bytes = reader.get_range(first_range).await.map_err(|source| {
1344 stream_copy_error("source read", source_path, destination_path, source)
1345 })?;
1346
1347 loop {
1348 let expected_bytes = current_range.len();
1349 if current_bytes.len() != expected_bytes {
1350 Span::current().record("validation", "failed");
1351 return Err(Error::io(format!(
1352 "multipart_stream_copy source range size mismatch from {source_path} to \
1353 {destination_path}: range={current_range:?}, \
1354 expected_bytes={expected_bytes}, actual_bytes={}",
1355 current_bytes.len()
1356 )));
1357 }
1358 bytes_transferred = bytes_transferred
1359 .checked_add(current_bytes.len())
1360 .ok_or_else(|| {
1361 Error::io(format!(
1362 "multipart_stream_copy byte count overflow from {source_path} to \
1363 {destination_path}"
1364 ))
1365 })?;
1366
1367 if bytes_transferred == source_size {
1368 writer.write_all(¤t_bytes).await.map_err(|source| {
1369 stream_copy_error(
1370 "destination write",
1371 source_path,
1372 destination_path,
1373 source,
1374 )
1375 })?;
1376 break;
1377 }
1378
1379 let range_end = bytes_transferred
1380 .checked_add(read_chunk_size)
1381 .unwrap_or(source_size)
1382 .min(source_size);
1383 let next_range = bytes_transferred..range_end;
1384 let next_read = reader.get_range(next_range.clone());
1385 let (write_result, next_bytes) =
1386 tokio::join!(writer.write_all(¤t_bytes), next_read);
1387 write_result.map_err(|source| {
1388 stream_copy_error("destination write", source_path, destination_path, source)
1389 })?;
1390 current_bytes = next_bytes.map_err(|source| {
1391 stream_copy_error("source read", source_path, destination_path, source)
1392 })?;
1393 current_range = next_range;
1394 }
1395 }
1396 Span::current().record("bytes_transferred", bytes_transferred as u64);
1397
1398 let write_result = Writer::shutdown(writer.as_mut()).await.map_err(|source| {
1399 stream_copy_error(
1400 "destination completion",
1401 source_path,
1402 destination_path,
1403 source,
1404 )
1405 })?;
1406 if write_result.size != source_size {
1407 Span::current().record("validation", "failed");
1408 return Err(Error::io(format!(
1409 "multipart_stream_copy writer size mismatch from {source_path} to \
1410 {destination_path}: source_size={source_size}, \
1411 writer_size={}",
1412 write_result.size
1413 )));
1414 }
1415
1416 let destination_size =
1417 destination_store
1418 .size(destination_path)
1419 .await
1420 .map_err(|source| {
1421 stream_copy_error(
1422 "destination validation",
1423 source_path,
1424 destination_path,
1425 source,
1426 )
1427 })?;
1428 Span::current().record("destination_size", destination_size);
1429 if destination_size != source_size as u64 {
1430 Span::current().record("validation", "failed");
1431 return Err(Error::io(format!(
1432 "multipart_stream_copy destination size mismatch from {source_path} to \
1433 {destination_path}: source_size={source_size}, \
1434 destination_size={destination_size}"
1435 )));
1436 }
1437
1438 Span::current().record("validation", "passed");
1439 Span::current().record("elapsed_ms", started_at.elapsed().as_millis() as u64);
1440 Ok(write_result)
1441 }
1442
1443 async fn copy_impl(
1449 &self,
1450 from: &Path,
1451 to: &Path,
1452 multipart_copy_fallback: bool,
1453 max_single_copy: u64,
1454 ) -> Result<()> {
1455 if self.has_direct_local_paths() {
1456 let metrics = self.io_tracker.begin_io("copy");
1458 let result = super::local::copy_file(from, to);
1459 metrics.record(&result, 0);
1460 return result;
1461 }
1462 if multipart_copy_fallback {
1463 let reader = self.open(from).await?;
1466 if reader.size().await? as u64 > max_single_copy {
1467 let mut writer = self.create(to).await?;
1468 writer.copy_from_reader(reader.as_ref()).await?;
1469 Writer::shutdown(writer.as_mut()).await?;
1470 return Ok(());
1471 }
1472 }
1473 Ok(self.inner.copy(from, to).await?)
1474 }
1475
1476 pub async fn read_dir(&self, dir_path: impl Into<Path>) -> Result<Vec<String>> {
1481 let path = dir_path.into();
1482 let path = Path::parse(&path)?;
1483 let output = self.inner.list_with_delimiter(Some(&path)).await?;
1484 Ok(output
1485 .common_prefixes
1486 .iter()
1487 .chain(output.objects.iter().map(|o| &o.location))
1488 .filter_map(|s| s.filename().map(|f| f.to_string()))
1489 .collect())
1490 }
1491
1492 pub async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
1499 Ok(self.inner.list_with_delimiter(prefix).await?)
1500 }
1501
1502 pub fn list(
1503 &self,
1504 path: Option<Path>,
1505 ) -> Pin<Box<dyn Stream<Item = Result<ObjectMeta>> + Send>> {
1506 Box::pin(ListRetryStream::new(self.inner.clone(), path, 5).map(|m| m.map_err(|e| e.into())))
1507 }
1508
1509 pub fn read_dir_all<'a, 'b>(
1513 &'a self,
1514 dir_path: impl Into<&'b Path> + Send,
1515 unmodified_since: Option<DateTime<Utc>>,
1516 ) -> BoxStream<'a, Result<ObjectMeta>> {
1517 self.inner.read_dir_all(dir_path, unmodified_since)
1518 }
1519
1520 pub async fn remove_dir_all(&self, dir_path: impl Into<Path>) -> Result<()> {
1522 let path = dir_path.into();
1523 let path = Path::parse(&path)?;
1524
1525 if let Some(local_dir_operations) = &self.local_dir_operations {
1526 let metrics = self.io_tracker.begin_io("delete");
1527 let result = local_dir_operations.remove_dir_all(&path).await;
1528 metrics.record(&result, 0);
1529 return result;
1530 }
1531 if self.has_direct_local_paths() {
1532 let metrics = self.io_tracker.begin_io("delete");
1536 let result = super::local::remove_dir_all(&path);
1537 metrics.record(&result, 0);
1538 return result;
1539 }
1540 let sub_entries = self
1541 .inner
1542 .list(Some(&path))
1543 .map(|m| m.map(|meta| meta.location))
1544 .boxed();
1545 self.inner
1546 .delete_stream(sub_entries)
1547 .try_collect::<Vec<_>>()
1548 .await?;
1549 if self.scheme == "file-object-store" {
1550 return super::local::remove_dir_all(&path);
1553 }
1554 Ok(())
1555 }
1556
1557 pub async fn remove_empty_dirs(
1583 &self,
1584 root_path: impl Into<Path>,
1585 retained_dirs: HashSet<Path>,
1586 verified_dirs: HashSet<Path>,
1587 unmodified_since: Option<DateTime<Utc>>,
1588 ) -> Result<()> {
1589 if !self.has_direct_local_paths() && self.scheme != "file-object-store" {
1590 return Ok(());
1591 }
1592
1593 let path = Path::parse(root_path.into())?;
1594 let metrics = self.io_tracker.begin_io("delete");
1595 let result = tokio::task::spawn_blocking(move || {
1596 super::local::remove_empty_dirs(&path, &retained_dirs, &verified_dirs, unmodified_since)
1597 })
1598 .await
1599 .map_err(|error| Error::io(format!("empty-directory cleanup task failed: {error}")))?;
1600 metrics.record(&result, 0);
1601 result
1602 }
1603
1604 pub fn remove_stream<'a>(
1619 &'a self,
1620 locations: BoxStream<'a, Result<Path>>,
1621 ) -> BoxStream<'a, Result<Path>> {
1622 let store = Arc::clone(&self.inner);
1623 locations
1624 .map(move |location| {
1625 let store = Arc::clone(&store);
1626 async move {
1627 let location = location?;
1628 match store.delete(&location).await {
1629 Ok(()) => Ok(location),
1630 Err(object_store::Error::NotFound { .. }) => Ok(location),
1631 Err(error) => Err(error.into()),
1632 }
1633 }
1634 })
1635 .buffer_unordered(self.io_parallelism())
1636 .boxed()
1637 }
1638
1639 pub async fn exists(&self, path: &Path) -> Result<bool> {
1641 match self.inner.head(path).await {
1642 Ok(_) => Ok(true),
1643 Err(object_store::Error::NotFound { path: _, source: _ }) => Ok(false),
1644 Err(e) => Err(e.into()),
1645 }
1646 }
1647
1648 pub async fn size(&self, path: &Path) -> Result<u64> {
1650 Ok(self.inner.head(path).await?.size)
1651 }
1652
1653 pub async fn read_one_all(&self, path: &Path) -> Result<Bytes> {
1655 let reader = self.open(path).await?;
1656 Ok(reader.get_all().await?)
1657 }
1658
1659 pub async fn read_one_range(&self, path: &Path, range: Range<usize>) -> Result<Bytes> {
1664 let reader = self.open(path).await?;
1665 Ok(reader.get_range(range).await?)
1666 }
1667}
1668
1669#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy)]
1671pub enum LanceConfigKey {
1672 DownloadRetryCount,
1674}
1675
1676impl FromStr for LanceConfigKey {
1677 type Err = Error;
1678
1679 fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
1680 match s.to_ascii_lowercase().as_str() {
1681 "download_retry_count" => Ok(Self::DownloadRetryCount),
1682 _ => Err(Error::invalid_input_source(
1683 format!("Invalid LanceConfigKey: {}", s).into(),
1684 )),
1685 }
1686 }
1687}
1688
1689#[derive(Clone, Debug, Default)]
1690pub struct StorageOptions(pub HashMap<String, String>);
1691
1692impl StorageOptions {
1693 pub fn new(options: HashMap<String, String>) -> Self {
1695 let mut options = options;
1696 if let Ok(value) = std::env::var("AZURE_STORAGE_ALLOW_HTTP") {
1697 options.insert("allow_http".into(), value);
1698 }
1699 if let Ok(value) = std::env::var("AZURE_STORAGE_USE_HTTP") {
1700 options.insert("allow_http".into(), value);
1701 }
1702 if let Ok(value) = std::env::var("AWS_ALLOW_HTTP") {
1703 options.insert("allow_http".into(), value);
1704 }
1705 if let Ok(value) = std::env::var("OBJECT_STORE_CLIENT_MAX_RETRIES") {
1706 options.insert("client_max_retries".into(), value);
1707 }
1708 if let Ok(value) = std::env::var("OBJECT_STORE_CLIENT_RETRY_TIMEOUT") {
1709 options.insert("client_retry_timeout".into(), value);
1710 }
1711 Self(options)
1712 }
1713
1714 pub fn allow_http(&self) -> bool {
1716 self.0.iter().any(|(key, value)| {
1717 key.to_ascii_lowercase().contains("allow_http") & str_is_truthy(value)
1718 })
1719 }
1720
1721 pub fn block_size(&self) -> Result<Option<usize>> {
1724 let Some((_, value)) = self
1725 .0
1726 .iter()
1727 .find(|(key, _)| key.eq_ignore_ascii_case("block_size"))
1728 else {
1729 return Ok(None);
1730 };
1731 let block_size = value.trim().parse::<usize>().map_err(|err| {
1732 Error::invalid_input(format!(
1733 "storage option block_size must be a number of bytes, got `{value}`: {err}"
1734 ))
1735 })?;
1736 Ok(Some(block_size))
1737 }
1738
1739 pub fn download_retry_count(&self) -> usize {
1741 self.0
1742 .iter()
1743 .find(|(key, _)| key.eq_ignore_ascii_case("download_retry_count"))
1744 .map(|(_, value)| value.parse::<usize>().unwrap_or(3))
1745 .unwrap_or(3)
1746 }
1747
1748 pub fn client_max_retries(&self) -> usize {
1750 self.0
1751 .iter()
1752 .find(|(key, _)| key.eq_ignore_ascii_case("client_max_retries"))
1753 .and_then(|(_, value)| value.parse::<usize>().ok())
1754 .unwrap_or(3)
1755 }
1756
1757 pub fn client_retry_timeout(&self) -> u64 {
1759 self.0
1760 .iter()
1761 .find(|(key, _)| key.eq_ignore_ascii_case("client_retry_timeout"))
1762 .and_then(|(_, value)| value.parse::<u64>().ok())
1763 .unwrap_or(180)
1764 }
1765
1766 pub fn get(&self, key: &str) -> Option<&String> {
1767 self.0.get(key)
1768 }
1769
1770 #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
1778 pub fn client_options(&self) -> Result<ClientOptions> {
1779 let mut headers = HeaderMap::new();
1780 for (key, value) in &self.0 {
1781 if let Some(header_name) = key.strip_prefix("headers.") {
1782 let name = header_name
1783 .parse::<http::header::HeaderName>()
1784 .map_err(|e| {
1785 Error::invalid_input(format!("invalid header name '{header_name}': {e}"))
1786 })?;
1787 let val = HeaderValue::from_str(value).map_err(|e| {
1788 Error::invalid_input(format!("invalid header value for '{header_name}': {e}"))
1789 })?;
1790 headers.insert(name, val);
1791 }
1792 }
1793 let mut client_options = ClientOptions::default();
1794 if !headers.is_empty() {
1795 client_options = client_options.with_default_headers(headers);
1796 }
1797 Ok(client_options)
1798 }
1799
1800 pub fn expires_at_millis(&self) -> Option<u64> {
1802 self.0
1803 .get(EXPIRES_AT_MILLIS_KEY)
1804 .and_then(|s| s.parse::<u64>().ok())
1805 }
1806}
1807
1808impl From<HashMap<String, String>> for StorageOptions {
1809 fn from(value: HashMap<String, String>) -> Self {
1810 Self::new(value)
1811 }
1812}
1813
1814static DEFAULT_OBJECT_STORE_REGISTRY: std::sync::LazyLock<ObjectStoreRegistry> =
1815 std::sync::LazyLock::new(ObjectStoreRegistry::default);
1816
1817impl ObjectStore {
1818 #[allow(clippy::too_many_arguments)]
1819 pub fn new(
1820 mut store: Arc<DynObjectStore>,
1821 location: Url,
1822 block_size: Option<usize>,
1823 wrapper: Option<Arc<dyn WrappingObjectStore>>,
1824 use_constant_size_upload_parts: bool,
1825 list_is_lexically_ordered: bool,
1826 io_parallelism: usize,
1827 download_retry_count: usize,
1828 storage_options: Option<&HashMap<String, String>>,
1829 ) -> Self {
1830 let scheme = location.scheme();
1831 let block_size = block_size.unwrap_or_else(|| infer_block_size(scheme));
1832 let store_prefix = match DEFAULT_OBJECT_STORE_REGISTRY.get_provider(scheme) {
1833 Some(provider) => provider
1834 .calculate_object_store_prefix(&location, storage_options)
1835 .unwrap(),
1836 None => {
1837 let store_prefix = format!("{}${}", location.scheme(), location.authority());
1838 log::warn!(
1839 "Guessing that object store prefix is {}, since object store scheme is not found in registry.",
1840 store_prefix
1841 );
1842 store_prefix
1843 }
1844 };
1845 let mut io_tracker = IOTracker::default();
1846 meter_store(
1847 &mut store,
1848 &mut io_tracker,
1849 &metrics_base(&store_prefix, &location),
1850 );
1851
1852 let store = match wrapper {
1853 Some(wrapper) => wrapper.wrap(&store_prefix, store),
1854 None => store,
1855 };
1856
1857 let tracked_store = io_tracker.wrap("", store);
1859
1860 Self {
1861 inner: tracked_store,
1862 local_dir_operations: None,
1863 scheme: scheme.into(),
1864 block_size,
1865 max_iop_size: *DEFAULT_MAX_IOP_SIZE,
1866 use_constant_size_upload_parts,
1867 list_is_lexically_ordered,
1868 io_parallelism,
1869 download_retry_count,
1870 io_tracker,
1871 store_prefix,
1872 paginated_lister: None,
1874 }
1875 }
1876}
1877
1878#[cfg(feature = "metrics")]
1887fn meter_store(inner: &mut Arc<dyn OSObjectStore>, io_tracker: &mut IOTracker, base: &str) {
1888 use crate::object_store::metrics::ObjectStoreMetricsExt;
1889 io_tracker.set_metrics_base(base);
1890 *inner = inner.clone().metered(base.to_owned());
1891}
1892
1893#[cfg(not(feature = "metrics"))]
1894fn meter_store(_inner: &mut Arc<dyn OSObjectStore>, _io_tracker: &mut IOTracker, _base: &str) {}
1895
1896#[cfg(feature = "metrics")]
1900pub(crate) fn metrics_base(store_prefix: &str, location: &Url) -> String {
1901 metrics::metrics_base(metrics::base_label_mode(), store_prefix, location)
1902}
1903
1904#[cfg(not(feature = "metrics"))]
1905pub(crate) fn metrics_base(store_prefix: &str, _location: &Url) -> String {
1906 store_prefix.to_owned()
1907}
1908
1909fn infer_block_size(scheme: &str) -> usize {
1910 match scheme {
1914 "file" => 4 * 1024,
1915 _ => 64 * 1024,
1916 }
1917}
1918
1919#[cfg(test)]
1920mod tests {
1921 use super::*;
1922 use async_trait::async_trait;
1923 use bytes::Bytes;
1924 use lance_core::utils::tempfile::{TempStdDir, TempStdFile, TempStrDir};
1925 use object_store::memory::InMemory;
1926 use object_store::{
1927 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOptions,
1928 PutOptions, PutPayload, PutResult, Result as OSResult, UploadPart,
1929 };
1930 use rstest::rstest;
1931 use serial_test::serial;
1932 use std::env::set_current_dir;
1933 use std::fmt::{Display, Formatter};
1934 use std::fs::{create_dir_all, write};
1935 use std::ops::Range;
1936 use std::path::Path as StdPath;
1937 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
1938
1939 fn write_to_file(path_str: &str, contents: &str) -> std::io::Result<()> {
1941 let path = expand_path(path_str).map_err(std::io::Error::other)?;
1942 std::fs::create_dir_all(path.parent().unwrap())?;
1943 write(path, contents)
1944 }
1945
1946 async fn read_from_store(store: &ObjectStore, path: &Path) -> Result<String> {
1947 let test_file_store = store.open(path).await.unwrap();
1948 let size = test_file_store.size().await.unwrap();
1949 let bytes = test_file_store.get_range(0..size).await.unwrap();
1950 let contents = String::from_utf8(bytes.to_vec()).unwrap();
1951 Ok(contents)
1952 }
1953
1954 #[tokio::test]
1955 async fn test_put_if_absent() {
1956 let temp_dir = TempStrDir::default();
1957 let path = Path::from(format!("{}/atomic-create", temp_dir.as_str()));
1958 let store = ObjectStore::local();
1959 store
1960 .put_if_absent(&path, Bytes::from_static(b"first").into())
1961 .await
1962 .unwrap();
1963 let error = store
1964 .put_if_absent(&path, Bytes::from_static(b"second").into())
1965 .await
1966 .unwrap_err();
1967 assert!(matches!(
1968 error,
1969 object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. }
1970 ));
1971 assert_eq!(
1972 store.read_one_all(&path).await.unwrap(),
1973 b"first".as_slice()
1974 );
1975 }
1976
1977 #[tokio::test]
1978 async fn test_put_if_absent_rejects_cos() {
1979 let mut store = ObjectStore::memory();
1980 store.scheme = "cos".to_string();
1981 let path = Path::from("atomic-create");
1982
1983 let error = store
1984 .put_if_absent(&path, Bytes::from_static(b"value").into())
1985 .await
1986 .unwrap_err();
1987
1988 assert!(matches!(error, object_store::Error::NotSupported { .. }));
1989 assert!(!store.exists(&path).await.unwrap());
1990 }
1991
1992 #[tokio::test]
1993 async fn test_io_parallelism_clamped_to_nonzero() {
1994 let store = ObjectStore::local();
1997 let mem_store = ObjectStore::memory();
2000 let path = Path::from("/io_parallelism_probe");
2001 mem_store.put(&path, b"x").await.unwrap();
2002
2003 unsafe { std::env::set_var("LANCE_IO_THREADS", "0") };
2006 assert_eq!(
2007 store.io_parallelism(),
2008 1,
2009 "LANCE_IO_THREADS=0 must clamp to 1"
2010 );
2011 assert_eq!(
2012 mem_store.open(&path).await.unwrap().io_parallelism(),
2013 1,
2014 "an opened reader must report the store's clamped parallelism"
2015 );
2016
2017 unsafe { std::env::set_var("LANCE_IO_THREADS", "8") };
2018 assert_eq!(
2019 store.io_parallelism(),
2020 8,
2021 "a positive override must pass through unchanged"
2022 );
2023 assert_eq!(
2024 mem_store.open(&path).await.unwrap().io_parallelism(),
2025 8,
2026 "an opened reader must honor the configured request limit"
2027 );
2028 assert_eq!(
2029 mem_store
2030 .open_with_size(&path, 1024 * 1024)
2031 .await
2032 .unwrap()
2033 .io_parallelism(),
2034 8,
2035 "a sized reader must honor the configured request limit"
2036 );
2037
2038 unsafe { std::env::remove_var("LANCE_IO_THREADS") };
2039 assert!(
2040 store.io_parallelism() >= 1,
2041 "the configured default parallelism must be at least 1"
2042 );
2043 }
2044
2045 #[tokio::test]
2046 async fn test_absolute_paths() {
2047 let tmp_path = TempStrDir::default();
2048 write_to_file(
2049 &format!("{tmp_path}/bar/foo.lance/test_file"),
2050 "TEST_CONTENT",
2051 )
2052 .unwrap();
2053
2054 for uri in &[
2056 format!("{tmp_path}/bar/foo.lance"),
2057 format!("{tmp_path}/./bar/foo.lance"),
2058 format!("{tmp_path}/bar/foo.lance/../foo.lance"),
2059 ] {
2060 let (store, path) = ObjectStore::from_uri(uri).await.unwrap();
2061 let contents = read_from_store(store.as_ref(), &path.clone().join("test_file"))
2062 .await
2063 .unwrap();
2064 assert_eq!(contents, "TEST_CONTENT");
2065 }
2066 }
2067
2068 #[tokio::test]
2069 async fn test_cloud_paths() {
2070 let uri = "s3://bucket/foo.lance";
2071 let (store, path) = ObjectStore::from_uri(uri).await.unwrap();
2072 assert_eq!(store.scheme, "s3");
2073 assert_eq!(path.to_string(), "foo.lance");
2074
2075 let (store, path) = ObjectStore::from_uri("s3+ddb://bucket/foo.lance")
2076 .await
2077 .unwrap();
2078 assert_eq!(store.scheme, "s3");
2079 assert_eq!(path.to_string(), "foo.lance");
2080
2081 let (store, path) = ObjectStore::from_uri("gs://bucket/foo.lance")
2082 .await
2083 .unwrap();
2084 assert_eq!(store.scheme, "gs");
2085 assert_eq!(path.to_string(), "foo.lance");
2086
2087 let (store, path) =
2088 ObjectStore::from_uri("abfss://filesystem@account.dfs.core.windows.net/foo.lance")
2089 .await
2090 .unwrap();
2091 assert_eq!(store.scheme, "abfss");
2092 assert_eq!(path.to_string(), "foo.lance");
2093 }
2094
2095 async fn test_block_size_used_test_helper(
2096 uri: &str,
2097 storage_options: Option<HashMap<String, String>>,
2098 default_expected_block_size: usize,
2099 ) {
2100 let registry = Arc::new(ObjectStoreRegistry::default());
2102 let accessor = storage_options
2103 .clone()
2104 .map(|opts| Arc::new(StorageOptionsAccessor::with_static_options(opts)));
2105 let params = ObjectStoreParams {
2106 storage_options_accessor: accessor.clone(),
2107 ..ObjectStoreParams::default()
2108 };
2109 let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms)
2110 .await
2111 .unwrap();
2112 assert_eq!(store.block_size, default_expected_block_size);
2113
2114 let registry = Arc::new(ObjectStoreRegistry::default());
2116 let params = ObjectStoreParams {
2117 block_size: Some(1024),
2118 storage_options_accessor: accessor,
2119 ..ObjectStoreParams::default()
2120 };
2121 let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms)
2122 .await
2123 .unwrap();
2124 assert_eq!(store.block_size, 1024);
2125
2126 let mut options_with_block_size = storage_options.unwrap_or_default();
2128 options_with_block_size.insert(String::from("block_size"), String::from("2048"));
2129 let accessor = Arc::new(StorageOptionsAccessor::with_static_options(
2130 options_with_block_size,
2131 ));
2132 let registry = Arc::new(ObjectStoreRegistry::default());
2133 let params = ObjectStoreParams {
2134 storage_options_accessor: Some(accessor.clone()),
2135 ..ObjectStoreParams::default()
2136 };
2137 let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms)
2138 .await
2139 .unwrap();
2140 assert_eq!(store.block_size, 2048);
2141
2142 let registry = Arc::new(ObjectStoreRegistry::default());
2144 let params = ObjectStoreParams {
2145 block_size: Some(1024),
2146 storage_options_accessor: Some(accessor),
2147 ..ObjectStoreParams::default()
2148 };
2149 let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms)
2150 .await
2151 .unwrap();
2152 assert_eq!(store.block_size, 1024);
2153 }
2154
2155 #[tokio::test]
2156 async fn test_block_size_option_rejects_invalid_values() {
2157 let registry = Arc::new(ObjectStoreRegistry::default());
2158 let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from(
2159 [(String::from("block_size"), String::from("64KiB"))],
2160 )));
2161 let params = ObjectStoreParams {
2162 storage_options_accessor: Some(accessor),
2163 ..ObjectStoreParams::default()
2164 };
2165 let error =
2166 ObjectStore::from_uri_and_params(registry, "memory:///bucket/foo.lance", ¶ms)
2167 .await
2168 .unwrap_err();
2169 assert!(
2170 matches!(error, lance_core::Error::InvalidInput { .. }),
2171 "{error:?}"
2172 );
2173 assert!(error.to_string().contains("block_size"), "{error}");
2174 }
2175
2176 #[rstest]
2177 #[case("s3://bucket/foo.lance", None)]
2178 #[case("gs://bucket/foo.lance", None)]
2179 #[case("az://account/bucket/foo.lance",
2180 Some(HashMap::from([
2181 (String::from("account_name"), String::from("account")),
2182 (String::from("container_name"), String::from("container"))
2183 ])))]
2184 #[case("abfss://filesystem@account.dfs.core.windows.net/foo.lance",
2185 Some(HashMap::from([
2186 (String::from("account_name"), String::from("account")),
2187 (String::from("container_name"), String::from("filesystem"))
2188 ])))]
2189 #[tokio::test]
2190 async fn test_block_size_used_cloud(
2191 #[case] uri: &str,
2192 #[case] storage_options: Option<HashMap<String, String>>,
2193 ) {
2194 test_block_size_used_test_helper(uri, storage_options, 64 * 1024).await;
2195 }
2196
2197 #[rstest]
2198 #[case("file")]
2199 #[case("file-object-store")]
2200 #[case("memory:///bucket/foo.lance")]
2201 #[tokio::test]
2202 async fn test_block_size_used_file(#[case] prefix: &str) {
2203 let tmp_path = TempStrDir::default();
2204 let path = format!("{tmp_path}/bar/foo.lance/test_file");
2205 write_to_file(&path, "URL").unwrap();
2206 let uri = format!("{prefix}:///{path}");
2207 test_block_size_used_test_helper(&uri, None, 4 * 1024).await;
2208 }
2209
2210 #[tokio::test]
2211 async fn test_relative_paths() {
2212 let tmp_path = TempStrDir::default();
2213 write_to_file(
2214 &format!("{tmp_path}/bar/foo.lance/test_file"),
2215 "RELATIVE_URL",
2216 )
2217 .unwrap();
2218
2219 set_current_dir(StdPath::new(tmp_path.as_ref())).expect("Error changing current dir");
2220 let (store, path) = ObjectStore::from_uri("./bar/foo.lance").await.unwrap();
2221
2222 let contents = read_from_store(store.as_ref(), &path.clone().join("test_file"))
2223 .await
2224 .unwrap();
2225 assert_eq!(contents, "RELATIVE_URL");
2226 }
2227
2228 #[tokio::test]
2229 async fn test_tilde_expansion() {
2230 let uri = "~/foo.lance";
2231 write_to_file(&format!("{uri}/test_file"), "TILDE").unwrap();
2232 let (store, path) = ObjectStore::from_uri(uri).await.unwrap();
2233 let contents = read_from_store(store.as_ref(), &path.clone().join("test_file"))
2234 .await
2235 .unwrap();
2236 assert_eq!(contents, "TILDE");
2237 }
2238
2239 #[tokio::test]
2240 async fn test_read_directory() {
2241 let path = TempStdDir::default();
2242 create_dir_all(path.join("foo").join("bar")).unwrap();
2243 create_dir_all(path.join("foo").join("zoo")).unwrap();
2244 create_dir_all(path.join("foo").join("zoo").join("abc")).unwrap();
2245 write_to_file(
2246 path.join("foo").join("test_file").to_str().unwrap(),
2247 "read_dir",
2248 )
2249 .unwrap();
2250 let (store, base) = ObjectStore::from_uri(path.to_str().unwrap()).await.unwrap();
2251
2252 let sub_dirs = store.read_dir(base.clone().join("foo")).await.unwrap();
2253 assert_eq!(sub_dirs, vec!["bar", "zoo", "test_file"]);
2254 }
2255
2256 #[tokio::test]
2257 async fn test_delete_directory_local_store() {
2258 test_delete_directory("").await;
2259 }
2260
2261 #[tokio::test]
2262 async fn test_delete_directory_file_object_store() {
2263 test_delete_directory("file-object-store").await;
2264 }
2265
2266 async fn test_delete_directory(scheme: &str) {
2267 let path = TempStdDir::default();
2268 create_dir_all(path.join("foo").join("bar")).unwrap();
2269 create_dir_all(path.join("foo").join("zoo")).unwrap();
2270 create_dir_all(path.join("foo").join("zoo").join("abc")).unwrap();
2271 write_to_file(
2272 path.join("foo")
2273 .join("bar")
2274 .join("test_file")
2275 .to_str()
2276 .unwrap(),
2277 "delete",
2278 )
2279 .unwrap();
2280 let file_url = Url::from_directory_path(&path).unwrap();
2281 let url = if scheme.is_empty() {
2282 file_url
2283 } else {
2284 let mut url = Url::parse(&format!("{scheme}:///")).unwrap();
2285 url.set_path(file_url.path());
2287 url
2288 };
2289 let (store, base) = ObjectStore::from_uri(url.as_ref()).await.unwrap();
2290 store
2291 .remove_dir_all(base.clone().join("foo"))
2292 .await
2293 .unwrap();
2294
2295 assert!(!path.join("foo").exists());
2296 }
2297
2298 #[rstest]
2299 #[case("file")]
2300 #[case("file-object-store")]
2301 #[tokio::test]
2302 async fn test_remove_empty_directories(#[case] scheme: &str) {
2303 let path = TempStdDir::default();
2304 let stale_dir = path.join("stale");
2305 let nested_stale_dir = path.join("nested_stale");
2306 let nested_stale_child = nested_stale_dir.join("child");
2307 create_dir_all(&stale_dir).unwrap();
2308 create_dir_all(&nested_stale_child).unwrap();
2309 create_dir_all(path.join("retained").join("child")).unwrap();
2310 write_to_file(
2311 path.join("file_bearing")
2312 .join("test_file")
2313 .to_str()
2314 .unwrap(),
2315 "keep",
2316 )
2317 .unwrap();
2318 create_dir_all(path.join("file_bearing").join("empty_child")).unwrap();
2319
2320 let file_url = Url::from_directory_path(&path).unwrap();
2321 let mut url = Url::parse(&format!("{scheme}:///")).unwrap();
2322 url.set_path(file_url.path());
2323 let (store, base) = ObjectStore::from_uri(url.as_ref()).await.unwrap();
2324
2325 #[cfg(unix)]
2326 let unmodified_since = {
2327 let old_modified_time =
2328 std::time::SystemTime::now() - std::time::Duration::from_secs(10 * 24 * 60 * 60);
2329 for directory in [&stale_dir, &nested_stale_dir, &nested_stale_child] {
2330 std::fs::File::open(directory)
2331 .unwrap()
2332 .set_times(std::fs::FileTimes::new().set_modified(old_modified_time))
2333 .unwrap();
2334 }
2335 DateTime::<Utc>::from(std::time::SystemTime::now())
2336 - chrono::TimeDelta::try_days(7).unwrap()
2337 };
2338 #[cfg(not(unix))]
2339 let unmodified_since = DateTime::<Utc>::from(std::time::SystemTime::now())
2340 + chrono::TimeDelta::try_days(1).unwrap();
2341
2342 store
2343 .remove_empty_dirs(
2344 base.clone(),
2345 HashSet::from([base.clone().join("retained")]),
2346 HashSet::new(),
2347 Some(unmodified_since),
2348 )
2349 .await
2350 .unwrap();
2351
2352 assert!(!path.join("stale").exists());
2353 assert!(!path.join("nested_stale").exists());
2354 assert!(path.join("retained").join("child").exists());
2355 assert!(path.join("file_bearing").join("empty_child").exists());
2356
2357 create_dir_all(path.join("fresh")).unwrap();
2358 create_dir_all(path.join("verified")).unwrap();
2359 store
2360 .remove_empty_dirs(
2361 base.clone(),
2362 HashSet::from([base.clone().join("retained")]),
2363 HashSet::from([base.clone().join("verified")]),
2364 Some(
2365 DateTime::<Utc>::from(std::time::SystemTime::now())
2366 - chrono::TimeDelta::try_days(7).unwrap(),
2367 ),
2368 )
2369 .await
2370 .unwrap();
2371
2372 assert!(path.join("fresh").exists());
2373 assert!(!path.join("verified").exists());
2374 }
2375
2376 #[derive(Debug)]
2377 struct TestWrapper {
2378 called: AtomicBool,
2379
2380 return_value: Arc<dyn OSObjectStore>,
2381 }
2382
2383 impl WrappingObjectStore for TestWrapper {
2384 fn wrap(
2385 &self,
2386 _store_prefix: &str,
2387 _original: Arc<dyn OSObjectStore>,
2388 ) -> Arc<dyn OSObjectStore> {
2389 self.called.store(true, Ordering::Relaxed);
2390
2391 self.return_value.clone()
2393 }
2394
2395 fn wrap_paginated(
2398 &self,
2399 _store_prefix: &str,
2400 _original: Arc<dyn PaginatedListStore>,
2401 ) -> Option<Arc<dyn PaginatedListStore>> {
2402 None
2403 }
2404 }
2405
2406 impl TestWrapper {
2407 fn called(&self) -> bool {
2408 self.called.load(Ordering::Relaxed)
2409 }
2410 }
2411
2412 #[derive(Debug)]
2414 struct StubLister;
2415
2416 #[async_trait]
2417 impl PaginatedListStore for StubLister {
2418 async fn list_paginated(
2419 &self,
2420 _prefix: Option<&str>,
2421 _opts: object_store::list::PaginatedListOptions,
2422 ) -> object_store::Result<object_store::list::PaginatedListResult> {
2423 unimplemented!("this lister exists to be wrapped, not to list")
2424 }
2425 }
2426
2427 #[derive(Debug)]
2429 struct PaginatedTestWrapper {
2430 name: &'static str,
2431 log: Arc<std::sync::Mutex<Vec<String>>>,
2432 }
2433
2434 impl WrappingObjectStore for PaginatedTestWrapper {
2435 fn wrap(
2436 &self,
2437 _store_prefix: &str,
2438 original: Arc<dyn OSObjectStore>,
2439 ) -> Arc<dyn OSObjectStore> {
2440 original
2441 }
2442
2443 fn wrap_paginated(
2444 &self,
2445 store_prefix: &str,
2446 original: Arc<dyn PaginatedListStore>,
2447 ) -> Option<Arc<dyn PaginatedListStore>> {
2448 self.log
2449 .lock()
2450 .unwrap()
2451 .push(format!("{}@{store_prefix}", self.name));
2452 Some(original)
2453 }
2454 }
2455
2456 #[rstest]
2460 #[case::every_wrapper_keeps_it(false, vec!["first@memory", "second@memory"])]
2461 #[case::one_wrapper_gives_it_up(true, vec!["first@memory"])]
2462 fn test_a_chain_wraps_the_lister_until_one_gives_it_up(
2463 #[case] gives_up: bool,
2464 #[case] expected_log: Vec<&str>,
2465 ) {
2466 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
2467 let mut wrappers: Vec<Arc<dyn WrappingObjectStore>> =
2468 vec![Arc::new(PaginatedTestWrapper {
2469 name: "first",
2470 log: log.clone(),
2471 })];
2472 if gives_up {
2473 wrappers.push(Arc::new(TestWrapper {
2474 called: AtomicBool::new(false),
2475 return_value: Arc::new(InMemory::new()),
2476 }));
2477 }
2478 wrappers.push(Arc::new(PaginatedTestWrapper {
2479 name: "second",
2480 log: log.clone(),
2481 }));
2482
2483 let wrapped = ChainedWrappingObjectStore::new(wrappers)
2484 .wrap_paginated("memory", Arc::new(StubLister));
2485
2486 assert_eq!(wrapped.is_none(), gives_up);
2487 assert_eq!(*log.lock().unwrap(), expected_log);
2488 }
2489
2490 #[rstest]
2494 #[case::gives_up_the_pushdown(true)]
2495 #[case::keeps_the_pushdown(false)]
2496 fn test_apply_wrapper_keeps_inner_and_the_lister_in_sync(#[case] gives_up: bool) {
2497 let replacement = Arc::new(InMemory::new());
2498 let giving_up = TestWrapper {
2499 called: AtomicBool::new(false),
2500 return_value: replacement.clone(),
2501 };
2502 let keeping = PaginatedTestWrapper {
2503 name: "passthrough",
2504 log: Arc::new(std::sync::Mutex::new(Vec::new())),
2505 };
2506 let wrapper: &dyn WrappingObjectStore = match gives_up {
2507 true => &giving_up,
2508 false => &keeping,
2509 };
2510
2511 let mut store = ObjectStore::memory();
2512 store.paginated_lister = Some(Arc::new(StubLister) as Arc<dyn PaginatedListStore>);
2513 store.apply_wrapper(wrapper);
2514
2515 assert_eq!(
2516 store.paginated_lister.is_some(),
2517 !gives_up,
2518 "the lister has to follow what the wrapper said"
2519 );
2520 assert_eq!(
2523 Arc::ptr_eq(&store.inner, &(replacement as Arc<dyn OSObjectStore>)),
2524 gives_up
2525 );
2526 }
2527
2528 #[tokio::test]
2529 async fn test_wrapper_identity_is_stable_across_tasks() {
2530 let wrapper = Arc::new(TestWrapper {
2531 called: AtomicBool::new(false),
2532 return_value: Arc::new(InMemory::new()),
2533 });
2534 let initial_params = ObjectStoreParams {
2535 object_store_wrapper: Some(wrapper.clone()),
2536 ..ObjectStoreParams::default()
2537 };
2538 let task_params = tokio::spawn(async move {
2539 ObjectStoreParams {
2540 object_store_wrapper: Some(wrapper),
2541 ..ObjectStoreParams::default()
2542 }
2543 })
2544 .await
2545 .unwrap();
2546
2547 assert_eq!(initial_params, task_params);
2548
2549 let mut initial_hasher = std::hash::DefaultHasher::new();
2550 std::hash::Hash::hash(&initial_params, &mut initial_hasher);
2551 let mut task_hasher = std::hash::DefaultHasher::new();
2552 std::hash::Hash::hash(&task_params, &mut task_hasher);
2553 assert_eq!(
2554 std::hash::Hasher::finish(&initial_hasher),
2555 std::hash::Hasher::finish(&task_hasher)
2556 );
2557 }
2558
2559 #[tokio::test]
2560 async fn test_wrapping_object_store_option_is_used() {
2561 let mock_inner_store: Arc<dyn OSObjectStore> = Arc::new(InMemory::new());
2563 let registry = Arc::new(ObjectStoreRegistry::default());
2564
2565 assert_eq!(Arc::strong_count(&mock_inner_store), 1);
2566
2567 let wrapper = Arc::new(TestWrapper {
2568 called: AtomicBool::new(false),
2569 return_value: mock_inner_store.clone(),
2570 });
2571
2572 let params = ObjectStoreParams {
2573 object_store_wrapper: Some(wrapper.clone()),
2574 ..ObjectStoreParams::default()
2575 };
2576
2577 assert!(!wrapper.called());
2579
2580 let _ = ObjectStore::from_uri_and_params(registry, "memory:///", ¶ms)
2581 .await
2582 .unwrap();
2583
2584 assert!(wrapper.called());
2586
2587 assert_eq!(Arc::strong_count(&mock_inner_store), 2);
2590 }
2591
2592 #[tokio::test]
2593 async fn test_local_paths() {
2594 let file_path = TempStdFile::default();
2595 let mut writer = ObjectStore::create_local_writer(&file_path).await.unwrap();
2596 writer.write_all(b"LOCAL").await.unwrap();
2597 Writer::shutdown(&mut writer).await.unwrap();
2598
2599 let reader = ObjectStore::open_local(&file_path).await.unwrap();
2600 let buf = reader.get_range(0..5).await.unwrap();
2601 assert_eq!(buf.as_ref(), b"LOCAL");
2602 }
2603
2604 #[cfg(unix)]
2605 #[tokio::test]
2606 async fn test_direct_local_writer_uses_standard_file_permissions() {
2607 let directory = TempStdDir::default();
2608 let reference_path = directory.join("reference");
2609 std::fs::File::create(&reference_path).unwrap();
2610 let expected_mode = std::fs::metadata(reference_path)
2611 .unwrap()
2612 .permissions()
2613 .mode()
2614 & 0o777;
2615
2616 let output_path = directory.join("output");
2617 let object_path = Path::from_absolute_path(&output_path).unwrap();
2618 let store = ObjectStore::local();
2619 let mut writer = store.create(&object_path).await.unwrap();
2620 writer.write_all(b"LOCAL").await.unwrap();
2621 Writer::shutdown(writer.as_mut()).await.unwrap();
2622
2623 let actual_mode = std::fs::metadata(output_path).unwrap().permissions().mode() & 0o777;
2624 assert_eq!(actual_mode, expected_mode);
2625 }
2626
2627 #[tokio::test]
2628 async fn test_read_one() {
2629 let file_path = TempStdFile::default();
2630 let mut writer = ObjectStore::create_local_writer(&file_path).await.unwrap();
2631 writer.write_all(b"LOCAL").await.unwrap();
2632 Writer::shutdown(&mut writer).await.unwrap();
2633
2634 let file_path_os = object_store::path::Path::parse(file_path.to_str().unwrap()).unwrap();
2635 let obj_store = ObjectStore::local();
2636 let buf = obj_store.read_one_all(&file_path_os).await.unwrap();
2637 assert_eq!(buf.as_ref(), b"LOCAL");
2638
2639 let buf = obj_store.read_one_range(&file_path_os, 0..5).await.unwrap();
2640 assert_eq!(buf.as_ref(), b"LOCAL");
2641 }
2642
2643 #[tokio::test]
2644 #[cfg(windows)]
2645 async fn test_windows_paths() {
2646 use std::path::Component;
2647 use std::path::Prefix;
2648 use std::path::Prefix::*;
2649
2650 fn get_path_prefix(path: &StdPath) -> Prefix<'_> {
2651 match path.components().next().unwrap() {
2652 Component::Prefix(prefix_component) => prefix_component.kind(),
2653 _ => panic!(),
2654 }
2655 }
2656
2657 fn get_drive_letter(prefix: Prefix) -> String {
2658 match prefix {
2659 Disk(bytes) => String::from_utf8(vec![bytes]).unwrap(),
2660 _ => panic!(),
2661 }
2662 }
2663
2664 let tmp_path = TempStdFile::default();
2665 let prefix = get_path_prefix(&tmp_path);
2666 let drive_letter = get_drive_letter(prefix);
2667
2668 write_to_file(
2669 &(format!("{drive_letter}:/test_folder/test.lance") + "/test_file"),
2670 "WINDOWS",
2671 )
2672 .unwrap();
2673
2674 for uri in &[
2675 format!("{drive_letter}:/test_folder/test.lance"),
2676 format!("{drive_letter}:\\test_folder\\test.lance"),
2677 ] {
2678 let (store, base) = ObjectStore::from_uri(uri).await.unwrap();
2679 let contents = read_from_store(store.as_ref(), &base.clone().join("test_file"))
2680 .await
2681 .unwrap();
2682 assert_eq!(contents, "WINDOWS");
2683 }
2684 }
2685
2686 #[tokio::test]
2687 async fn test_cross_filesystem_copy() {
2688 let source_dir = TempStdDir::default();
2690 let dest_dir = TempStdDir::default();
2691
2692 let source_file_name = "test_file.txt";
2694 let source_file = source_dir.join(source_file_name);
2695 std::fs::write(&source_file, b"test content").unwrap();
2696
2697 let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap())
2699 .await
2700 .unwrap();
2701
2702 let from_path = base_path.clone().join(source_file_name);
2704
2705 let dest_file = dest_dir.join("copied_file.txt");
2707 let dest_str = dest_file.to_str().unwrap();
2708 let to_path = object_store::path::Path::parse(dest_str).unwrap();
2709
2710 store.copy(&from_path, &to_path).await.unwrap();
2712
2713 assert!(dest_file.exists());
2715 let copied_content = std::fs::read(&dest_file).unwrap();
2716 assert_eq!(copied_content, b"test content");
2717 }
2718
2719 #[tokio::test]
2720 async fn test_copy_creates_parent_directories() {
2721 let source_dir = TempStdDir::default();
2722 let dest_dir = TempStdDir::default();
2723
2724 let source_file_name = "test_file.txt";
2726 let source_file = source_dir.join(source_file_name);
2727 std::fs::write(&source_file, b"test content").unwrap();
2728
2729 let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap())
2731 .await
2732 .unwrap();
2733
2734 let from_path = base_path.clone().join(source_file_name);
2736
2737 let dest_file = dest_dir.join("nested").join("dirs").join("copied_file.txt");
2739 let dest_str = dest_file.to_str().unwrap();
2740 let to_path = object_store::path::Path::parse(dest_str).unwrap();
2741
2742 store.copy(&from_path, &to_path).await.unwrap();
2744
2745 assert!(dest_file.exists());
2747 assert!(dest_file.parent().unwrap().exists());
2748 let copied_content = std::fs::read(&dest_file).unwrap();
2749 assert_eq!(copied_content, b"test content");
2750 }
2751
2752 #[derive(Debug)]
2757 struct CopyFailingStore {
2758 inner: InMemory,
2759 }
2760
2761 impl Display for CopyFailingStore {
2762 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
2763 write!(f, "CopyFailingStore")
2764 }
2765 }
2766
2767 #[derive(Debug, Default)]
2768 struct MultipartObservations {
2769 part_count: AtomicUsize,
2770 abort_count: AtomicUsize,
2771 native_copy_count: AtomicUsize,
2772 delete_stream_count: AtomicUsize,
2773 delete_in_flight: AtomicUsize,
2774 delete_max_in_flight: AtomicUsize,
2776 }
2777
2778 struct InFlightGuard(Arc<MultipartObservations>);
2780
2781 impl Drop for InFlightGuard {
2782 fn drop(&mut self) {
2783 self.0.delete_in_flight.fetch_sub(1, Ordering::SeqCst);
2784 }
2785 }
2786
2787 #[derive(Debug)]
2788 struct ObservedMultipartUpload {
2789 inner: Box<dyn MultipartUpload>,
2790 observations: Arc<MultipartObservations>,
2791 fail_parts: bool,
2792 }
2793
2794 #[async_trait]
2795 impl MultipartUpload for ObservedMultipartUpload {
2796 fn put_part(&mut self, data: PutPayload) -> UploadPart {
2797 self.observations.part_count.fetch_add(1, Ordering::SeqCst);
2798 if self.fail_parts {
2799 return Box::pin(async {
2800 Err(object_store::Error::Generic {
2801 store: "ObservedMultipartStore",
2802 source: "injected multipart part failure".into(),
2803 })
2804 });
2805 }
2806 self.inner.put_part(data)
2807 }
2808
2809 async fn complete(&mut self) -> OSResult<PutResult> {
2810 self.inner.complete().await
2811 }
2812
2813 async fn abort(&mut self) -> OSResult<()> {
2814 self.observations.abort_count.fetch_add(1, Ordering::SeqCst);
2815 self.inner.abort().await
2816 }
2817 }
2818
2819 #[derive(Debug)]
2820 struct ObservedMultipartStore {
2821 inner: InMemory,
2822 observations: Arc<MultipartObservations>,
2823 fail_parts: bool,
2824 destination_size_adjustment: u64,
2825 }
2826
2827 impl Display for ObservedMultipartStore {
2828 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
2829 write!(f, "ObservedMultipartStore")
2830 }
2831 }
2832
2833 #[async_trait]
2834 impl OSObjectStore for ObservedMultipartStore {
2835 async fn put_opts(
2836 &self,
2837 location: &Path,
2838 bytes: PutPayload,
2839 opts: PutOptions,
2840 ) -> OSResult<PutResult> {
2841 self.inner.put_opts(location, bytes, opts).await
2842 }
2843
2844 async fn put_multipart_opts(
2845 &self,
2846 location: &Path,
2847 opts: PutMultipartOptions,
2848 ) -> OSResult<Box<dyn MultipartUpload>> {
2849 let inner = self.inner.put_multipart_opts(location, opts).await?;
2850 Ok(Box::new(ObservedMultipartUpload {
2851 inner,
2852 observations: self.observations.clone(),
2853 fail_parts: self.fail_parts,
2854 }))
2855 }
2856
2857 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
2858 let is_head = options.head;
2859 let mut result = self.inner.get_opts(location, options).await?;
2860 if is_head && location.filename() == Some("destination.bin") {
2861 result.meta.size = result
2862 .meta
2863 .size
2864 .checked_add(self.destination_size_adjustment)
2865 .expect("test destination size should not overflow");
2866 }
2867 Ok(result)
2868 }
2869
2870 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
2871 self.inner.get_ranges(location, ranges).await
2872 }
2873
2874 fn delete_stream(
2875 &self,
2876 locations: BoxStream<'static, OSResult<Path>>,
2877 ) -> BoxStream<'static, OSResult<Path>> {
2878 let observations = self.observations.clone();
2879 observations
2880 .delete_stream_count
2881 .fetch_add(1, Ordering::SeqCst);
2882 let in_flight = observations.delete_in_flight.fetch_add(1, Ordering::SeqCst) + 1;
2883 observations
2884 .delete_max_in_flight
2885 .fetch_max(in_flight, Ordering::SeqCst);
2886 let guard = InFlightGuard(observations);
2887 let inner = self.inner.delete_stream(locations);
2888 async move {
2891 tokio::task::yield_now().await;
2892 tokio::time::sleep(Duration::from_millis(5)).await;
2893 inner.map(move |result| {
2894 let _ = &guard;
2895 result
2896 })
2897 }
2898 .flatten_stream()
2899 .boxed()
2900 }
2901
2902 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
2903 self.inner.list(prefix)
2904 }
2905
2906 fn list_with_offset(
2907 &self,
2908 prefix: Option<&Path>,
2909 offset: &Path,
2910 ) -> BoxStream<'static, OSResult<ObjectMeta>> {
2911 self.inner.list_with_offset(prefix, offset)
2912 }
2913
2914 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
2915 self.inner.list_with_delimiter(prefix).await
2916 }
2917
2918 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
2919 self.observations
2920 .native_copy_count
2921 .fetch_add(1, Ordering::SeqCst);
2922 self.inner.copy_opts(from, to, opts).await
2923 }
2924 }
2925
2926 #[async_trait]
2927 impl OSObjectStore for CopyFailingStore {
2928 async fn put_opts(
2929 &self,
2930 location: &Path,
2931 bytes: PutPayload,
2932 opts: PutOptions,
2933 ) -> OSResult<PutResult> {
2934 self.inner.put_opts(location, bytes, opts).await
2935 }
2936 async fn put_multipart_opts(
2937 &self,
2938 location: &Path,
2939 opts: PutMultipartOptions,
2940 ) -> OSResult<Box<dyn MultipartUpload>> {
2941 self.inner.put_multipart_opts(location, opts).await
2942 }
2943 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
2944 self.inner.get_opts(location, options).await
2945 }
2946 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
2947 self.inner.get_ranges(location, ranges).await
2948 }
2949 fn delete_stream(
2950 &self,
2951 locations: BoxStream<'static, OSResult<Path>>,
2952 ) -> BoxStream<'static, OSResult<Path>> {
2953 self.inner.delete_stream(locations)
2954 }
2955 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
2956 self.inner.list(prefix)
2957 }
2958 fn list_with_offset(
2959 &self,
2960 prefix: Option<&Path>,
2961 offset: &Path,
2962 ) -> BoxStream<'static, OSResult<ObjectMeta>> {
2963 self.inner.list_with_offset(prefix, offset)
2964 }
2965 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
2966 self.inner.list_with_delimiter(prefix).await
2967 }
2968 async fn copy_opts(&self, _from: &Path, _to: &Path, _opts: CopyOptions) -> OSResult<()> {
2969 Err(object_store::Error::Generic {
2970 store: "CopyFailingStore",
2971 source: "single-shot copy disabled in test".into(),
2972 })
2973 }
2974 }
2975
2976 #[tokio::test]
2977 async fn test_copy_streams_objects_larger_than_threshold() {
2978 let mut store = ObjectStore::memory();
2984 store.inner = Arc::new(CopyFailingStore {
2985 inner: InMemory::new(),
2986 });
2987
2988 let from = Path::from("source.bin");
2989 let contents = b"streaming multipart copy payload well past the tiny threshold";
2990 store.put(&from, contents).await.unwrap();
2991
2992 let streamed = Path::from("streamed.bin");
2995 store.copy_impl(&from, &streamed, true, 8).await.unwrap();
2996 let copied = store.read_one_all(&streamed).await.unwrap();
2997 assert_eq!(copied.as_ref(), contents.as_slice());
2998
2999 let native = Path::from("native.bin");
3003 assert!(
3004 store
3005 .copy_impl(&from, &native, true, u64::MAX)
3006 .await
3007 .is_err()
3008 );
3009 }
3010
3011 #[tokio::test]
3012 async fn test_copy_via_stream_never_uses_native_copy() {
3013 let mut store = ObjectStore::memory();
3014 store.inner = Arc::new(CopyFailingStore {
3015 inner: InMemory::new(),
3016 });
3017
3018 let source = Path::from("source.bin");
3019 let destination = Path::from("destination.bin");
3020 let contents = b"stream raw bytes instead of issuing native copy";
3021 store.put(&source, contents).await.unwrap();
3022
3023 let result = store
3024 .copy_via_stream(&source, &store, &destination)
3025 .await
3026 .unwrap();
3027
3028 assert_eq!(result.size, contents.len());
3029 assert_eq!(
3030 store.read_one_all(&destination).await.unwrap().as_ref(),
3031 contents
3032 );
3033 }
3034
3035 #[tokio::test]
3036 async fn test_remove_stream_removes_every_path() {
3037 let store = ObjectStore::memory();
3038 let paths: Vec<Path> = (0..64).map(|i| Path::from(format!("obj-{i:03}"))).collect();
3039 for path in &paths {
3040 store.put(path, b"x").await.unwrap();
3041 }
3042
3043 let to_remove = futures::stream::iter(paths.clone().into_iter().map(Ok)).boxed();
3044 let mut reported = store
3045 .remove_stream(to_remove)
3046 .try_collect::<Vec<_>>()
3047 .await
3048 .unwrap();
3049
3050 reported.sort_unstable();
3052 let mut expected = paths.clone();
3053 expected.sort_unstable();
3054 assert_eq!(reported, expected);
3055 for path in &paths {
3056 assert!(!store.exists(path).await.unwrap(), "{path} still present");
3057 }
3058 }
3059
3060 #[tokio::test]
3061 async fn test_remove_stream_tolerates_already_deleted() {
3062 let store = ObjectStore::memory();
3065 let present = Path::from("present");
3066 let absent = Path::from("never-written");
3067 store.put(&present, b"x").await.unwrap();
3068
3069 let to_remove =
3070 futures::stream::iter(vec![Ok(absent.clone()), Ok(present.clone())]).boxed();
3071 let mut removed = store
3072 .remove_stream(to_remove)
3073 .try_collect::<Vec<_>>()
3074 .await
3075 .expect("an already-absent path must not fail the stream");
3076
3077 removed.sort_unstable();
3078 let mut expected = vec![absent, present.clone()];
3079 expected.sort_unstable();
3080 assert_eq!(removed, expected);
3081 assert!(!store.exists(&present).await.unwrap());
3082 }
3083
3084 #[tokio::test]
3085 async fn test_remove_stream_deletes_concurrently() {
3086 let observations = Arc::new(MultipartObservations::default());
3089 let mut store = ObjectStore::memory();
3090 store.inner = Arc::new(ObservedMultipartStore {
3091 inner: InMemory::new(),
3092 observations: observations.clone(),
3093 fail_parts: false,
3094 destination_size_adjustment: 0,
3095 });
3096
3097 let paths: Vec<Path> = (0..64).map(|i| Path::from(format!("obj-{i:03}"))).collect();
3098 for path in &paths {
3099 store.put(path, b"x").await.unwrap();
3100 }
3101
3102 let to_remove = futures::stream::iter(paths.clone().into_iter().map(Ok)).boxed();
3103 let removed = store
3104 .remove_stream(to_remove)
3105 .try_collect::<Vec<_>>()
3106 .await
3107 .unwrap();
3108
3109 assert_eq!(removed.len(), paths.len());
3110 assert!(
3111 observations.delete_max_in_flight.load(Ordering::SeqCst) > 1,
3112 "deletions ran one at a time; remove_stream must overlap them"
3113 );
3114 }
3115
3116 #[tokio::test]
3117 async fn test_bulk_copy_streams_when_server_side_copy_is_disabled() {
3118 let observations = Arc::new(MultipartObservations::default());
3119 let mut store = ObjectStore::memory();
3120 store.inner = Arc::new(ObservedMultipartStore {
3121 inner: InMemory::new(),
3122 observations: observations.clone(),
3123 fail_parts: false,
3124 destination_size_adjustment: 0,
3125 });
3126
3127 let source = Path::from("source.bin");
3128 let destination = Path::from("destination.bin");
3129 let contents = b"stream by default";
3130 store.put(&source, contents).await.unwrap();
3131
3132 let result = store
3133 .copy_bulk_with_server_side_copy(&source, &store, &destination, false)
3134 .await
3135 .unwrap();
3136
3137 assert_eq!(result.size, contents.len());
3138 assert_eq!(observations.native_copy_count.load(Ordering::SeqCst), 0);
3139 assert_eq!(
3140 store.read_one_all(&destination).await.unwrap().as_ref(),
3141 contents
3142 );
3143 }
3144
3145 #[test]
3146 #[serial(server_side_copy_env)]
3147 fn test_server_side_copy_environment_policy() {
3148 let previous_value = std::env::var_os(SERVER_SIDE_COPY_ENABLED_ENV);
3149 let mut store = ObjectStore::memory();
3150 store.scheme = "test-cloud".to_string();
3151 let destination_store = store.clone();
3152
3153 unsafe { std::env::remove_var(SERVER_SIDE_COPY_ENABLED_ENV) };
3156 assert!(!store.uses_server_side_copy(&destination_store));
3157
3158 unsafe { std::env::set_var(SERVER_SIDE_COPY_ENABLED_ENV, "true") };
3160 assert!(store.uses_server_side_copy(&destination_store));
3161
3162 unsafe {
3164 match previous_value {
3165 Some(value) => std::env::set_var(SERVER_SIDE_COPY_ENABLED_ENV, value),
3166 None => std::env::remove_var(SERVER_SIDE_COPY_ENABLED_ENV),
3167 }
3168 }
3169 }
3170
3171 #[tokio::test]
3172 async fn test_bulk_copy_uses_server_side_copy_when_enabled_for_same_store() {
3173 let observations = Arc::new(MultipartObservations::default());
3174 let mut source_store = ObjectStore::memory();
3175 source_store.scheme = "test-cloud".to_string();
3176 source_store.inner = Arc::new(ObservedMultipartStore {
3177 inner: InMemory::new(),
3178 observations: observations.clone(),
3179 fail_parts: false,
3180 destination_size_adjustment: 0,
3181 });
3182 let destination_store = source_store.clone();
3183
3184 let source = Path::from("source.bin");
3185 let destination = Path::from("destination.bin");
3186 let contents = b"use native copy when explicitly enabled";
3187 source_store.put(&source, contents).await.unwrap();
3188
3189 let result = source_store
3190 .copy_bulk_with_server_side_copy(&source, &destination_store, &destination, true)
3191 .await
3192 .unwrap();
3193
3194 assert_eq!(result.size, contents.len());
3195 assert_eq!(observations.native_copy_count.load(Ordering::SeqCst), 1);
3196 assert_eq!(
3197 destination_store
3198 .read_one_all(&destination)
3199 .await
3200 .unwrap()
3201 .as_ref(),
3202 contents
3203 );
3204 }
3205
3206 #[tokio::test]
3207 async fn test_bulk_copy_streams_for_distinct_clients_with_same_prefix() {
3208 let shared_inner = InMemory::new();
3209 let source_observations = Arc::new(MultipartObservations::default());
3210 let mut source_store = ObjectStore::memory();
3211 source_store.scheme = "test-cloud".to_string();
3212 source_store.store_prefix = "test-cloud$bucket".to_string();
3213 source_store.inner = Arc::new(ObservedMultipartStore {
3214 inner: shared_inner.clone(),
3215 observations: source_observations.clone(),
3216 fail_parts: false,
3217 destination_size_adjustment: 0,
3218 });
3219 let destination_observations = Arc::new(MultipartObservations::default());
3220 let mut destination_store = ObjectStore::memory();
3221 destination_store.scheme = "test-cloud".to_string();
3222 destination_store.store_prefix = "test-cloud$bucket".to_string();
3223 destination_store.inner = Arc::new(ObservedMultipartStore {
3224 inner: shared_inner,
3225 observations: destination_observations.clone(),
3226 fail_parts: false,
3227 destination_size_adjustment: 0,
3228 });
3229
3230 let source = Path::from("source.bin");
3231 let destination = Path::from("destination.bin");
3232 let contents = b"use native copy when explicitly enabled";
3233 source_store.put(&source, contents).await.unwrap();
3234
3235 let result = source_store
3236 .copy_bulk_with_server_side_copy(&source, &destination_store, &destination, true)
3237 .await
3238 .unwrap();
3239
3240 assert_eq!(result.size, contents.len());
3241 assert_eq!(
3242 source_observations.native_copy_count.load(Ordering::SeqCst),
3243 0
3244 );
3245 assert_eq!(
3246 destination_observations
3247 .native_copy_count
3248 .load(Ordering::SeqCst),
3249 0
3250 );
3251 assert_eq!(
3252 destination_store
3253 .read_one_all(&destination)
3254 .await
3255 .unwrap()
3256 .as_ref(),
3257 contents
3258 );
3259 }
3260
3261 #[tokio::test]
3262 async fn test_bulk_copy_rejects_server_side_destination_size_mismatch() {
3263 let observations = Arc::new(MultipartObservations::default());
3264 let mut source_store = ObjectStore::memory();
3265 source_store.scheme = "test-cloud".to_string();
3266 source_store.inner = Arc::new(ObservedMultipartStore {
3267 inner: InMemory::new(),
3268 observations: observations.clone(),
3269 fail_parts: false,
3270 destination_size_adjustment: 1,
3271 });
3272 let destination_store = source_store.clone();
3273
3274 let source = Path::from("source.bin");
3275 let destination = Path::from("destination.bin");
3276 source_store
3277 .put(&source, b"validate native copy")
3278 .await
3279 .unwrap();
3280
3281 let error = source_store
3282 .copy_bulk_with_server_side_copy(&source, &destination_store, &destination, true)
3283 .await
3284 .unwrap_err();
3285
3286 assert!(
3287 error.to_string().contains("destination size mismatch"),
3288 "expected validation failure, got: {error}"
3289 );
3290 assert_eq!(observations.native_copy_count.load(Ordering::SeqCst), 1);
3291 }
3292
3293 #[tokio::test]
3294 async fn test_bulk_copy_streams_across_stores_when_server_side_copy_is_enabled() {
3295 let source_store = ObjectStore::memory();
3296 let observations = Arc::new(MultipartObservations::default());
3297 let mut destination_store = ObjectStore::memory();
3298 destination_store.inner = Arc::new(ObservedMultipartStore {
3299 inner: InMemory::new(),
3300 observations: observations.clone(),
3301 fail_parts: false,
3302 destination_size_adjustment: 0,
3303 });
3304
3305 let source = Path::from("source.bin");
3306 let destination = Path::from("destination.bin");
3307 let contents = b"cross-store copies must stream";
3308 source_store.put(&source, contents).await.unwrap();
3309
3310 let result = source_store
3311 .copy_bulk_with_server_side_copy(&source, &destination_store, &destination, true)
3312 .await
3313 .unwrap();
3314
3315 assert_eq!(result.size, contents.len());
3316 assert_eq!(observations.native_copy_count.load(Ordering::SeqCst), 0);
3317 assert_eq!(
3318 destination_store
3319 .read_one_all(&destination)
3320 .await
3321 .unwrap()
3322 .as_ref(),
3323 contents
3324 );
3325 }
3326
3327 #[tokio::test]
3328 async fn test_copy_via_stream_preserves_local_not_found() {
3329 let directory = TempStdDir::default();
3330 let (store, base_path) = ObjectStore::from_uri(directory.to_str().unwrap())
3331 .await
3332 .unwrap();
3333 let source = base_path.clone().join("missing.bin");
3334 let destination = base_path.join("destination.bin");
3335
3336 let error = store
3337 .copy_via_stream(&source, &store, &destination)
3338 .await
3339 .unwrap_err();
3340
3341 assert!(
3342 error.is_not_found(),
3343 "expected not-found error, got: {error}"
3344 );
3345 }
3346
3347 #[tokio::test]
3348 async fn test_copy_via_stream_uses_multiple_parts() {
3349 let mut source_store = ObjectStore::memory();
3350 source_store.max_iop_size = 1024 * 1024;
3351 let observations = Arc::new(MultipartObservations::default());
3352 let mut destination_store = ObjectStore::memory();
3353 destination_store.inner = Arc::new(ObservedMultipartStore {
3354 inner: InMemory::new(),
3355 observations: observations.clone(),
3356 fail_parts: false,
3357 destination_size_adjustment: 0,
3358 });
3359
3360 let source = Path::from("source.bin");
3361 let destination = Path::from("destination.bin");
3362 let contents = vec![42; crate::object_writer::initial_upload_size() * 2 + 1];
3363 source_store.put(&source, &contents).await.unwrap();
3364
3365 let result = source_store
3366 .copy_via_stream(&source, &destination_store, &destination)
3367 .await
3368 .unwrap();
3369
3370 assert_eq!(result.size, contents.len());
3371 assert!(
3372 observations.part_count.load(Ordering::SeqCst) >= 2,
3373 "stream copy should split a large destination into multiple upload parts"
3374 );
3375 assert_eq!(
3376 destination_store
3377 .read_one_all(&destination)
3378 .await
3379 .unwrap()
3380 .as_ref(),
3381 contents.as_slice()
3382 );
3383 }
3384
3385 #[tokio::test]
3386 async fn test_copy_via_stream_aborts_failed_upload_and_retains_source() {
3387 let source_store = ObjectStore::memory();
3388 let observations = Arc::new(MultipartObservations::default());
3389 let mut destination_store = ObjectStore::memory();
3390 destination_store.inner = Arc::new(ObservedMultipartStore {
3391 inner: InMemory::new(),
3392 observations: observations.clone(),
3393 fail_parts: true,
3394 destination_size_adjustment: 0,
3395 });
3396
3397 let source = Path::from("source.bin");
3398 let destination = Path::from("destination.bin");
3399 let contents = vec![7; crate::object_writer::initial_upload_size() * 2];
3400 source_store.put(&source, &contents).await.unwrap();
3401
3402 let error = source_store
3403 .copy_via_stream(&source, &destination_store, &destination)
3404 .await
3405 .unwrap_err();
3406 let error_message = error.to_string();
3407 assert!(
3408 (error_message.contains("destination write")
3409 || error_message.contains("destination completion"))
3410 && error_message.contains("injected multipart part failure"),
3411 "expected upload-stage context and the underlying error, got: {error}"
3412 );
3413
3414 tokio::time::timeout(Duration::from_secs(1), async {
3415 loop {
3416 if observations.abort_count.load(Ordering::SeqCst) > 0 {
3417 break;
3418 }
3419 tokio::task::yield_now().await;
3420 }
3421 })
3422 .await
3423 .expect("multipart abort should complete");
3424 assert_eq!(observations.abort_count.load(Ordering::SeqCst), 1);
3425 assert_eq!(
3426 source_store.read_one_all(&source).await.unwrap().as_ref(),
3427 contents.as_slice()
3428 );
3429 assert!(!destination_store.exists(&destination).await.unwrap());
3430 }
3431
3432 #[tokio::test]
3433 async fn test_copy_via_stream_rejects_destination_size_mismatch() {
3434 let source_store = ObjectStore::memory();
3435 let mut destination_store = ObjectStore::memory();
3436 destination_store.inner = Arc::new(ObservedMultipartStore {
3437 inner: InMemory::new(),
3438 observations: Arc::new(MultipartObservations::default()),
3439 fail_parts: false,
3440 destination_size_adjustment: 1,
3441 });
3442
3443 let source = Path::from("source.bin");
3444 let destination = Path::from("destination.bin");
3445 let contents = b"validate the destination after completion";
3446 source_store.put(&source, contents).await.unwrap();
3447
3448 let error = source_store
3449 .copy_via_stream(&source, &destination_store, &destination)
3450 .await
3451 .unwrap_err();
3452
3453 assert!(
3454 error.to_string().contains("destination size mismatch"),
3455 "expected validation failure, got: {error}"
3456 );
3457 }
3458
3459 #[test]
3460 #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
3461 fn test_client_options_extracts_headers() {
3462 let opts = StorageOptions(HashMap::from([
3463 ("headers.x-custom-foo".to_string(), "bar".to_string()),
3464 ("headers.x-ms-version".to_string(), "2023-11-03".to_string()),
3465 ("region".to_string(), "us-west-2".to_string()),
3466 ]));
3467 let client_options = opts.client_options().unwrap();
3468
3469 let opts_no_headers = StorageOptions(HashMap::from([(
3472 "region".to_string(),
3473 "us-west-2".to_string(),
3474 )]));
3475 opts_no_headers.client_options().unwrap();
3476
3477 #[cfg(feature = "gcp")]
3481 {
3482 use object_store::gcp::GoogleCloudStorageBuilder;
3483 let _builder = GoogleCloudStorageBuilder::new()
3484 .with_client_options(client_options)
3485 .with_url("gs://test-bucket");
3486 }
3487 }
3488
3489 #[test]
3490 #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
3491 fn test_client_options_rejects_invalid_header_name() {
3492 let opts = StorageOptions(HashMap::from([(
3493 "headers.bad header".to_string(),
3494 "value".to_string(),
3495 )]));
3496 let err = opts.client_options().unwrap_err();
3497 assert!(err.to_string().contains("invalid header name"));
3498 }
3499
3500 #[test]
3501 #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
3502 fn test_client_options_rejects_invalid_header_value() {
3503 let opts = StorageOptions(HashMap::from([(
3504 "headers.x-good-name".to_string(),
3505 "bad\x01value".to_string(),
3506 )]));
3507 let err = opts.client_options().unwrap_err();
3508 assert!(err.to_string().contains("invalid header value"));
3509 }
3510
3511 #[test]
3512 #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))]
3513 fn test_client_options_empty_when_no_header_keys() {
3514 let opts = StorageOptions(HashMap::from([
3515 ("region".to_string(), "us-east-1".to_string()),
3516 ("access_key_id".to_string(), "AKID".to_string()),
3517 ]));
3518 opts.client_options().unwrap();
3519 }
3520}