Skip to main content

lance_io/
object_store.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Extend [object_store::ObjectStore] functionalities
5
6use 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
77// Local disks tend to do fine with a few threads
78// Note: the number of threads here also impacts the number of files
79// we need to read in some situations.  So keeping this at 8 keeps the
80// RAM on our scanner down.
81pub const DEFAULT_LOCAL_IO_PARALLELISM: usize = 8;
82// Cloud disks often need many many threads to saturate the network
83pub 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; // 4KB block size
88#[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; // 64KB block size
99
100pub 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    /// Returns true if the file exists.
158    async fn exists(&self, path: &Path) -> Result<bool>;
159
160    /// Read all files (start from base directory) recursively
161    ///
162    /// unmodified_since can be specified to only return files that have not been modified since the given time.
163    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/// Wraps [ObjectStore](object_store::ObjectStore)
202#[derive(Clone)]
203pub struct ObjectStore {
204    // Inner object store
205    pub inner: Arc<dyn OSObjectStore>,
206    // Provider-owned native directory operations for rooted local stores.
207    local_dir_operations: Option<Arc<dyn LocalDirOperations>>,
208    scheme: String,
209    block_size: usize,
210    max_iop_size: u64,
211    /// Whether to use constant size upload parts for multipart uploads. This
212    /// is only necessary for Cloudflare R2.
213    pub use_constant_size_upload_parts: bool,
214    /// Whether we can assume that the list of files is lexically ordered. This
215    /// is true for object stores, but not for local filesystems.
216    pub list_is_lexically_ordered: bool,
217    io_parallelism: usize,
218    /// Number of times to retry a failed download
219    download_retry_count: usize,
220    /// IO tracker for monitoring read/write operations
221    io_tracker: IOTracker,
222    /// The datastore prefix that uniquely identifies this object store. It encodes information
223    /// which usually cannot be found in the URL such as Azure account name. The prefix plus the
224    /// path uniquely identifies any object inside the store.
225    pub store_prefix: String,
226    /// The backend's paginated listing API, when it has one. `None` means
227    /// [`Self::read_dir_page`] has to list a directory in full to page through it.
228    pub(crate) paginated_lister: Option<Arc<dyn PaginatedListStore>>,
229}
230
231// Hand-written because `PaginatedListStore` is not `Debug`.
232impl 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        // We aren't counting `inner` here which is problematic but an ObjectStore
256        // shouldn't be too big.  The only exception might be the write cache but, if
257        // the writer cache has data, it means we're using it somewhere else that isn't
258        // a cache and so that doesn't really count.
259        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    /// Wrap an object store with additional functionality
271    ///
272    /// The store_prefix is a string which uniquely identifies the object
273    /// store being wrapped.
274    fn wrap(&self, store_prefix: &str, original: Arc<dyn OSObjectStore>) -> Arc<dyn OSObjectStore>;
275
276    /// Wrap the paginated listing API that goes with the store, if it has one.
277    ///
278    /// [`ObjectStore::read_dir_page`] pushes the page size and the resume position into
279    /// [`PaginatedListStore`], which is a separate trait from [`OSObjectStore`] and so cannot
280    /// be reached through the store [`Self::wrap`] returns. A listing that is pushed down
281    /// therefore does not pass through [`Self::wrap`], and this is where a wrapper says what
282    /// should happen instead:
283    ///
284    /// - `Some(lister)` keeps the pushdown, wrapping the lister or handing back the one
285    ///   given. Right for a wrapper that observes rather than intercepts — metering, caching,
286    ///   mirroring writes.
287    /// - `None` gives up the pushdown, so listings go through [`Self::wrap`] as a full
288    ///   directory read. Right for a wrapper that hides, rewrites or fails paths, which a
289    ///   pushed-down listing would otherwise walk straight past.
290    ///
291    /// A wrapper that keeps the pushdown must leave the listing itself alone: setting
292    /// [`offset`](object_store::list::PaginatedListOptions::offset) or changing the delimiter
293    /// breaks paging, since `read_dir_page` reads one directory level and resumes by the token
294    /// it got back.
295    ///
296    /// There is deliberately no default: getting this wrong is either a silent loss of speed
297    /// or a silent loss of the wrapper, and neither announces itself.
298    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    // One wrapper giving up the pushdown gives it up for the chain: the listing has to go
328    // through `wrap`, which is every wrapper in the chain at once.
329    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/// Parameters to create an [ObjectStore]
341///
342#[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    /// Refresh offset for AWS credentials when using the legacy AWS credentials path.
348    /// For StorageOptionsAccessor, use `refresh_offset_millis` storage option instead.
349    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    /// Unified storage options accessor with caching and automatic refresh
354    ///
355    /// Provides storage options and optionally a dynamic provider for automatic
356    /// credential refresh. Use `StorageOptionsAccessor::with_static_options()` for static
357    /// options or `StorageOptionsAccessor::with_initial_and_provider()` for dynamic refresh.
358    pub storage_options_accessor: Option<Arc<StorageOptionsAccessor>>,
359    /// Use constant size upload parts for multipart uploads. Only necessary
360    /// for Cloudflare R2, which doesn't support variable size parts. When this
361    /// is false, max upload size is 2.5TB. When this is true, the max size is
362    /// 50GB.
363    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    /// Get the StorageOptionsAccessor from the params
386    pub fn get_accessor(&self) -> Option<Arc<StorageOptionsAccessor>> {
387        self.storage_options_accessor.clone()
388    }
389
390    /// Get storage options from the accessor, if any
391    ///
392    /// Returns the initial storage options from the accessor without triggering refresh.
393    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    /// The block size to use: the explicit `block_size` parameter, else the
400    /// `block_size` storage option, else `None` for the store's default.
401    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    /// Resolve these params for a single base path scope.
412    ///
413    /// Storage options may carry base-scoped entries (`base_<id>.<key>`) that
414    /// apply only to one registered base path; see
415    /// [`StorageOptionsAccessor::scoped_to_base`]. Returns the params unchanged
416    /// when the storage options contain no base-scoped entries.
417    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    // Trait object pointers include vtable metadata, which is not stable across codegen units.
435    // Cache identity must follow the Arc allocation instead.
436    Arc::as_ptr(wrapper) as *const ()
437}
438
439// We implement hash for caching
440impl std::hash::Hash for ObjectStoreParams {
441    #[allow(deprecated)]
442    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
443        // For hashing, we use pointer values for ObjectStore, S3 credentials, wrapper
444        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
465// We implement eq for caching
466impl 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        // For equality, we use pointer comparison for ObjectStore, S3 credentials, wrapper
476        // For accessor, we use accessor_id() for semantic equality
477        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
508/// Convert a URI string or local path to a URL
509///
510/// This function handles both proper URIs (with schemes like `file://`, `s3://`, etc.)
511/// and plain local filesystem paths. On Windows, it correctly handles drive letters
512/// that might be parsed as URL schemes.
513///
514/// # Examples
515///
516/// ```
517/// # use lance_io::object_store::uri_to_url;
518/// // URIs are preserved
519/// let url = uri_to_url("s3://bucket/path").unwrap();
520/// assert_eq!(url.scheme(), "s3");
521///
522/// // Local paths are converted to file:// URIs
523/// # #[cfg(unix)]
524/// let url = uri_to_url("/tmp/data").unwrap();
525/// # #[cfg(unix)]
526/// assert_eq!(url.scheme(), "file");
527/// ```
528pub fn uri_to_url(uri: &str) -> Result<Url> {
529    match Url::parse(uri) {
530        Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
531            // On Windows, the drive is parsed as a scheme
532            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    // path_abs::PathAbs::new(".") returns an empty string.
548    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    // Accept forms with repo type prefix (models/datasets/spaces) or legacy without.
584    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    /// Parse from a string URI.
622    ///
623    /// Returns the ObjectStore instance and the absolute path to the object.
624    ///
625    /// This uses the default [ObjectStoreRegistry] to find the object store. To
626    /// allow for potential re-use of object store instances, it's recommended to
627    /// create a shared [ObjectStoreRegistry] and pass that to [Self::from_uri_and_params].
628    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    /// Parse from a string URI.
635    ///
636    /// Returns the ObjectStore instance and the absolute path to the object.
637    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    /// Parse a URI and build a fresh object store outside the registry cache.
646    ///
647    /// The caller must retain the returned store for as long as its
648    /// provider-local state should be reused.
649    #[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            // Always wrap with IO tracking
682            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                // Type-erased on the way in, so there is no telling if it can paginate.
697                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        // We know the scheme is valid if we got a store back.
710        let provider = registry.get_provider(url.scheme()).expect_ok()?;
711        let path = provider.extract_path(&url)?;
712
713        Ok((store, path))
714    }
715
716    /// Extract the path component from a URI without initializing the object store.
717    ///
718    /// This is a synchronous operation that only parses the URI and extracts the path,
719    /// without creating or initializing any object store instance.
720    ///
721    /// # Arguments
722    ///
723    /// * `registry` - The object store registry to get the provider
724    /// * `uri` - The URI to extract the path from
725    ///
726    /// # Returns
727    ///
728    /// The extracted path component
729    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    /// Local object store.
749    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    /// Create a in-memory object store directly for testing.
759    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    /// Returns true if the object store pointed to a local file system.
769    pub fn is_local(&self) -> bool {
770        self.scheme == "file" || self.scheme == "file+uring"
771    }
772
773    /// Returns true when object paths directly encode absolute local filesystem paths.
774    ///
775    /// Local stores rooted below the filesystem root, such as UNC-backed stores, use
776    /// their inner object-store implementation instead of direct filesystem access.
777    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    /// Whether this object store prefers the lite scheduler.
789    ///
790    /// The lite scheduler is designed for backends like io_uring where
791    /// tasks should only be polled when the consumer polls them.
792    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    /// The amount of parallelism to use for I/O operations.
809    ///
810    /// Honors the `LANCE_IO_THREADS` override when set, otherwise the store's configured value.
811    /// Always at least 1: callers feed this straight into `buffered` / `buffer_unordered`, and a
812    /// window of 0 makes those streams never poll their input — e.g. a metadata-only `count_rows`
813    /// would hang rather than return.
814    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    /// Get the IO tracker for this object store
822    ///
823    /// The IO tracker can be used to get statistics about read/write operations
824    /// performed on this object store.
825    pub fn io_tracker(&self) -> &IOTracker {
826        &self.io_tracker
827    }
828
829    /// Get a snapshot of current IO statistics without resetting counters
830    ///
831    /// Returns the current IO statistics without modifying the internal state.
832    /// Use this when you need to check stats without resetting them.
833    pub fn io_stats_snapshot(&self) -> IoStats {
834        self.io_tracker.stats()
835    }
836
837    /// Get incremental IO statistics since the last call to this method
838    ///
839    /// Returns the accumulated statistics since the last call and resets the
840    /// counters to zero. This is useful for tracking IO operations between
841    /// different stages of processing.
842    pub fn io_stats_incremental(&self) -> IoStats {
843        self.io_tracker.incremental_stats()
844    }
845
846    /// Apply a [`WrappingObjectStore`] to both `inner` and `paginated_lister` together.
847    ///
848    /// Keeps both halves in sync: a wrapper returning `None` from
849    /// [`WrappingObjectStore::wrap_paginated`] clears the lister so that
850    /// [`Self::read_dir_page`] falls back through the (already-wrapped) `inner`.
851    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    /// Open a file for path.
860    ///
861    /// Parameters
862    /// - ``path``: Absolute path to the file.
863    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                // Check if current-thread mode enabled
877                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    /// Open a reader for a file with known size.
913    ///
914    /// This size may either have been retrieved from a list operation or
915    /// cached metadata. By passing in the known size, we can skip a HEAD / metadata
916    /// call.
917    pub async fn open_with_size(&self, path: &Path, known_size: usize) -> Result<Box<dyn Reader>> {
918        // If we know the file is really small, we can read the whole thing
919        // as a single request.
920        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                // Check if current-thread mode enabled
942                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    /// Create an [ObjectWriter] from local [std::path::Path]
978    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    /// Open an [Reader] from local [std::path::Path]
986    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    /// Create a new file.
994    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                        // NamedTempFile defaults to 0o600. Use ordinary file creation permissions so the published file honors the caller's umask.
1010                        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    /// A helper function to create a file and write content to it.
1033    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    /// Atomically creates an object without replacing an existing object.
1040    ///
1041    /// Local stores publish a uniquely named staging object with a conditional
1042    /// rename. Other stores use their conditional create operation. Tencent COS
1043    /// is rejected because it can silently ignore conditional create requests.
1044    ///
1045    /// Returns [`object_store::Error::NotSupported`] without writing when the
1046    /// backend cannot reliably provide put-if-absent semantics.
1047    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    /// AWS S3 and GCS reject a single-shot server-side copy whose source is
1096    /// larger than this; such sources are streamed through a multipart write.
1097    const MAX_SINGLE_COPY_BYTES: u64 = 5 * 1024 * 1024 * 1024; // 5 GiB
1098
1099    pub async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
1100        // S3 and GCS cap single-shot server-side copies at 5 GiB and object_store
1101        // does not fall back to a multipart copy for larger sources
1102        // (https://github.com/apache/arrow-rs-object-store/issues/563). Azure and
1103        // other blob stores don't have this limit, so we only pay for the fallback
1104        // (an extra size lookup) on S3 and GCS.
1105        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    /// Copy an object using the policy for bulk file movement.
1116    ///
1117    /// Streaming is the default because it works across object stores and does
1118    /// not require provider-native copy support. Setting
1119    /// `LANCE_IO_SERVER_SIDE_COPY_ENABLED` to a truthy value opts same-store
1120    /// copies into [`Self::copy`]. Cross-store and local copies continue to use
1121    /// [`Self::copy_via_stream`].
1122    ///
1123    /// ```no_run
1124    /// # use lance_core::Result;
1125    /// # use lance_io::object_store::ObjectStore;
1126    /// # use object_store::path::Path;
1127    /// # async fn copy(source: &ObjectStore, destination: &ObjectStore) -> Result<()> {
1128    /// source
1129    ///     .copy_bulk(
1130    ///         &Path::from("staging/index.lance"),
1131    ///         destination,
1132    ///         &Path::from("index.lance"),
1133    ///     )
1134    ///     .await?;
1135    /// # Ok(())
1136    /// # }
1137    /// ```
1138    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        // Prefixes can collide across endpoints or wrappers, where native copy could
1198        // read or write the wrong backend. Exact client identity is required.
1199        self.is_cloud()
1200            && destination_store.is_cloud()
1201            && Arc::ptr_eq(&self.inner, &destination_store.inner)
1202    }
1203
1204    /// Copy an object by streaming its bytes through Lance's multipart-aware writer.
1205    ///
1206    /// Unlike [`Self::copy`], this never delegates to a provider-native server-side
1207    /// copy. The source and destination may use different object stores. The copy
1208    /// succeeds only after the byte count reported by the writer and a destination
1209    /// metadata lookup both match the source size.
1210    ///
1211    /// ```no_run
1212    /// # use lance_core::Result;
1213    /// # use lance_io::object_store::ObjectStore;
1214    /// # use object_store::path::Path;
1215    /// # async fn copy(source: &ObjectStore, destination: &ObjectStore) -> Result<()> {
1216    /// source
1217    ///     .copy_via_stream(
1218    ///         &Path::from("staging/index.lance"),
1219    ///         destination,
1220    ///         &Path::from("index.lance"),
1221    ///     )
1222    ///     .await?;
1223    /// # Ok(())
1224    /// # }
1225    /// ```
1226    #[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(&current_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(&current_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    /// Copy `from` to `to`. When `multipart_copy_fallback` is set, a source
1444    /// larger than `max_single_copy` is streamed through a multipart write
1445    /// instead of a single-shot server-side copy. Both are parameters so tests
1446    /// can drive the streaming path without a multi-gigabyte fixture or an S3
1447    /// endpoint.
1448    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            // Use std::fs::copy for local filesystem to support cross-filesystem copies
1457            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            // Reuse the reader for both the size lookup (a single cached HEAD)
1464            // and the streamed copy, avoiding a separate HEAD request.
1465            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    /// Read a directory (start from base directory) and returns all sub-paths in the directory.
1477    ///
1478    /// This enumerates the whole prefix before it returns, however many children it holds.
1479    /// Use [`Self::read_dir_page`] to page through a directory instead.
1480    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    /// Non-recursive, path-segment delimited list of a single directory level.
1493    ///
1494    /// Unlike [`Self::list`], which recurses into the entire subtree, this returns
1495    /// only the immediate children of `prefix`: the child "directories" as
1496    /// [`ListResult::common_prefixes`] and the direct child files as
1497    /// [`ListResult::objects`].
1498    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    /// Read all files (start from base directory) recursively
1510    ///
1511    /// unmodified_since can be specified to only return files that have not been modified since the given time.
1512    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    /// Remove a directory recursively.
1521    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            // The local file system provider needs to delete both files and directories.
1533            // Counted as a single delete request, matching how `delete_stream`
1534            // counts one batched request regardless of how many paths it removes.
1535            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            // file-object-store tries to do everything as similarly as possible to the remote
1551            // object stores. But we still have to delete the directory entries afterwards.
1552            return super::local::remove_dir_all(&path);
1553        }
1554        Ok(())
1555    }
1556
1557    /// Remove eligible materialized empty directories below a local root.
1558    ///
1559    /// This is a no-op for object stores, which do not materialize directories.
1560    /// Traversal does not follow symbolic links. Directories in `retained_dirs` and their
1561    /// descendants are preserved. Other directories are removed only if they are empty and
1562    /// either appear in `verified_dirs` or predate `unmodified_since`. Passing `None` for
1563    /// `unmodified_since` disables the age check.
1564    ///
1565    /// ```
1566    /// # use std::collections::HashSet;
1567    /// # use chrono::Utc;
1568    /// # use lance_core::Result;
1569    /// # use lance_io::object_store::ObjectStore;
1570    /// # async fn remove_stale_index_dirs(store: &ObjectStore) -> Result<()> {
1571    /// store
1572    ///     .remove_empty_dirs(
1573    ///         "dataset/_indices",
1574    ///         HashSet::new(),
1575    ///         HashSet::new(),
1576    ///         Some(Utc::now()),
1577    ///     )
1578    ///     .await?;
1579    /// # Ok(())
1580    /// # }
1581    /// ```
1582    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    /// Delete every location in the stream, up to `io_parallelism` at a time.
1605    ///
1606    /// Deleting serially costs a round trip per object — ~30/s against a cloud store,
1607    /// which is days for a table with millions of unreferenced files.
1608    ///
1609    /// `ObjectStore::delete_stream` would also batch (S3: 1000 keys per request) but
1610    /// needs a `'static` stream of `object_store::Result`; callers pass a borrowed
1611    /// stream of `lance_core::Result`. Using it means changing this signature.
1612    ///
1613    /// Order is not preserved; no caller may depend on it.
1614    ///
1615    /// A location that is already gone counts as removed. Callers list first and
1616    /// delete after, so a concurrent writer or a second cleanup can remove a path in
1617    /// between; failing there would abandon an entire sweep over one absent object.
1618    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    /// Check a file exists.
1640    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    /// Get file size.
1649    pub async fn size(&self, path: &Path) -> Result<u64> {
1650        Ok(self.inner.head(path).await?.size)
1651    }
1652
1653    /// Convenience function to open a reader and read all the bytes
1654    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    /// Convenience function open a reader and make a single request
1660    ///
1661    /// If you will be making multiple requests to the path it is more efficient to call [`Self::open`]
1662    /// and then call [`Reader::get_range`] multiple times.
1663    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/// Options that can be set for multiple object stores
1670#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy)]
1671pub enum LanceConfigKey {
1672    /// Number of times to retry a download that fails
1673    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    /// Create a new instance of [`StorageOptions`]
1694    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    /// Denotes if unsecure connections via http are allowed
1715    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    /// Byte gap below which the I/O scheduler merges two reads of one file
1722    /// into a single request, overriding the store's default.
1723    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    /// Number of times to retry a download that fails
1740    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    /// Max retry times to set in RetryConfig for object store client
1749    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    /// Seconds of timeout to set in RetryConfig for object store client
1758    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    /// Build [`ClientOptions`] with default headers extracted from `headers.*` keys.
1771    ///
1772    /// Keys prefixed with `headers.` are parsed into HTTP headers. For example,
1773    /// `headers.x-ms-version = 2023-11-03` results in a default header
1774    /// `x-ms-version: 2023-11-03`.
1775    ///
1776    /// Returns an error if any `headers.*` key has an invalid header name or value.
1777    #[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    /// Get the expiration time in milliseconds since epoch, if present
1801    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        // Always wrap with IO tracking
1858        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            // Type-erased on the way in, so there is no telling if it can paginate.
1873            paginated_lister: None,
1874        }
1875    }
1876}
1877
1878/// Wrap `inner` so its operations publish metrics labelled by `store_prefix`,
1879/// and label `io_tracker` with the same prefix so the local reads and writes
1880/// that bypass `inner` publish under it too.
1881///
1882/// The two go together on purpose: a store metered on one path but not the other
1883/// would report a partial picture that reads like a complete one. Every
1884/// constructor that hands an [`ObjectStore`] to a caller must route its `inner`
1885/// through here, or through nothing at all.
1886#[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/// The `base` metrics label for a store opened at `location`; see
1897/// [`metrics::metrics_base`]. Without the `metrics` feature the label is unused
1898/// and the prefix keeps the registry cache keyed per bucket as before.
1899#[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    // Block size: On local file systems, we use 4KB block size. On cloud
1911    // object stores, we use 64KB block size. This is generally the largest
1912    // block size where we don't see a latency penalty.
1913    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    /// Write test content to file.
1940    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        // `io_parallelism()` feeds `buffered`/`buffer_unordered` windows; a value of 0 makes those
1995        // streams never poll, hanging callers (e.g. a metadata-only `count_rows`). It must clamp.
1996        let store = ObjectStore::local();
1997        // Readers opened by the store must advertise the store's normalized
1998        // effective parallelism, not the hardcoded cloud default.
1999        let mem_store = ObjectStore::memory();
2000        let path = Path::from("/io_parallelism_probe");
2001        mem_store.put(&path, b"x").await.unwrap();
2002
2003        // SAFETY: process-global env var, set and restored within this test. `io_parallelism()`
2004        // only reads it, and a concurrent reader observes a valid clamped value, never 0.
2005        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        // test a few variations of the same path
2055        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        // Test the default
2101        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, &params)
2110            .await
2111            .unwrap();
2112        assert_eq!(store.block_size, default_expected_block_size);
2113
2114        // Ensure param is used
2115        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, &params)
2122            .await
2123            .unwrap();
2124        assert_eq!(store.block_size, 1024);
2125
2126        // The storage option applies when no parameter is given...
2127        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, &params)
2138            .await
2139            .unwrap();
2140        assert_eq!(store.block_size, 2048);
2141
2142        // ...and the explicit parameter wins over it.
2143        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, &params)
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", &params)
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            // Use the file:// URL's normalized path so this works on Windows too.
2286            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            // return a mocked value so we can check if the final store is the one we expect
2392            self.return_value.clone()
2393        }
2394
2395        // This one swaps the store out entirely, so a listing that went around it would be
2396        // listing something else.
2397        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    /// A lister that exists only to be wrapped.
2413    #[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    /// Records the listers it was handed, and leaves the store alone.
2428    #[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    /// A chain hands the lister to each of its wrappers in turn. One wrapper giving up the
2457    /// pushdown gives it up for the chain, and the wrappers after it are never asked: the
2458    /// listing is going through `wrap` either way, which is every wrapper at once.
2459    #[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    /// `apply_wrapper` keeps both halves of the store in sync. A wrapper that gives up the
2491    /// pushdown has to clear the lister too, or `read_dir_page` would keep talking to the
2492    /// backend behind the wrapper's back.
2493    #[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        // The wrapper that gives up the pushdown is also the one that swaps the store out, so
2521        // whether `inner` was replaced says that `wrap` ran on the same wrapper.
2522        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        // Make a store for the inner store first
2562        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        // not called yet
2578        assert!(!wrapper.called());
2579
2580        let _ = ObjectStore::from_uri_and_params(registry, "memory:///", &params)
2581            .await
2582            .unwrap();
2583
2584        // called after construction
2585        assert!(wrapper.called());
2586
2587        // hard to compare two trait pointers as the point to vtables
2588        // using the ref count as a proxy to make sure that the store is correctly kept
2589        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        // Create two temporary directories that simulate different filesystems
2689        let source_dir = TempStdDir::default();
2690        let dest_dir = TempStdDir::default();
2691
2692        // Create a test file in the source directory
2693        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        // Create ObjectStore for local filesystem
2698        let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap())
2699            .await
2700            .unwrap();
2701
2702        // Create paths relative to the ObjectStore base
2703        let from_path = base_path.clone().join(source_file_name);
2704
2705        // Use object_store::Path::parse for the destination
2706        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        // Perform the copy operation
2711        store.copy(&from_path, &to_path).await.unwrap();
2712
2713        // Verify the file was copied correctly
2714        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        // Create a test file in the source directory
2725        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        // Create ObjectStore for local filesystem
2730        let (store, base_path) = ObjectStore::from_uri(source_dir.to_str().unwrap())
2731            .await
2732            .unwrap();
2733
2734        // Create paths
2735        let from_path = base_path.clone().join(source_file_name);
2736
2737        // Create destination with nested directories that don't exist yet
2738        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        // Perform the copy operation - should create parent directories
2743        store.copy(&from_path, &to_path).await.unwrap();
2744
2745        // Verify the file was copied correctly and directories were created
2746        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    /// Inner store that forwards everything to `InMemory` except single-shot
2753    /// server-side copy (`copy_opts`), which always fails. This lets a test
2754    /// prove that `ObjectStore::copy` fell back to a streaming multipart copy
2755    /// for an oversized source rather than issuing a single `CopyObject`.
2756    #[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        /// High-water mark of concurrent deletions; 1 means serial.
2775        delete_max_in_flight: AtomicUsize,
2776    }
2777
2778    /// Decrements the in-flight count when its stream is dropped.
2779    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            // An in-memory delete resolves on first poll, so overlap is only
2889            // observable if the deletion yields.
2890            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        // memory:// is non-local but isn't an S3/GCS scheme, so copy() wouldn't
2979        // enable the fallback on its own. Drive copy_impl directly with
2980        // multipart_copy_fallback = true to exercise the streaming path. The
2981        // inner store rejects any single-shot copy, so a successful copy can only
2982        // have gone through the streaming branch.
2983        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        // Source size (61 bytes) exceeds the threshold -> must stream via a
2993        // multipart write rather than a single-shot server-side copy.
2994        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        // Source size below the threshold -> single-shot copy, which the inner
3000        // store rejects, confirming that the streaming branch (not native copy)
3001        // is what made the first copy succeed.
3002        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        // Concurrent deletion does not preserve order.
3051        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        // Callers list first and delete after, so a path can disappear in between.
3063        // Failing there abandons the whole sweep over one absent object.
3064        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        // Serial deletion is a round trip per object; the high-water mark is 1
3087        // exactly when that regression happens.
3088        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        // SAFETY: this serialized test is the only test that mutates this task-specific
3154        // environment variable, and it restores the original value before returning.
3155        unsafe { std::env::remove_var(SERVER_SIDE_COPY_ENABLED_ENV) };
3156        assert!(!store.uses_server_side_copy(&destination_store));
3157
3158        // SAFETY: see the serialized-test guarantee above.
3159        unsafe { std::env::set_var(SERVER_SIDE_COPY_ENABLED_ENV, "true") };
3160        assert!(store.uses_server_side_copy(&destination_store));
3161
3162        // SAFETY: restore the process environment before the test returns.
3163        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        // Verify non-header keys are not consumed as headers by creating
3470        // another StorageOptions with no headers.* keys.
3471        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        // Smoke test: the client_options with headers should be usable
3478        // in a builder (we can't inspect the headers directly, but building
3479        // should not fail).
3480        #[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}