diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 4490647a5..a455a5fae 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -16,8 +16,15 @@ use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore}; use lance_datafusion::utils::StreamingWriteSource; use lance_encoding::version::LanceFileVersion; use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; -use lance_table::io::commit::commit_handler_from_url; +use lance_table::{ + format::{IndexMetadata, Manifest, Transaction as LanceTableTransaction}, + io::commit::{ + CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, ManifestWriter, + commit_handler_from_url, + }, +}; use object_store::local::LocalFileSystem; +use object_store::path::Path as ObjectPath; use snafu::ResultExt; use crate::blob::{ensure_blob_storage_version, has_blob_columns}; @@ -25,6 +32,8 @@ use crate::connection::ConnectRequest; use crate::database::ReadConsistency; use crate::database::namespace::LanceNamespaceDatabase; use crate::error::{CreateDirSnafu, Error, Result}; +#[cfg(windows)] +use crate::io::object_store::rooted_file_commit_handler; use crate::io::object_store::{MirroringObjectStoreWrapper, new_default_session}; use crate::table::NativeTable; use crate::utils::validate_table_name; @@ -56,6 +65,71 @@ fn session_or_default( } } +/// Routes clone source metadata through the source commit handler while all +/// destination operations continue to use the target handler. +#[derive(Debug)] +struct CloneCommitHandler { + source_base: ObjectPath, + source: Arc, + target: Arc, +} + +#[async_trait::async_trait] +impl CommitHandler for CloneCommitHandler { + async fn resolve_latest_location( + &self, + base_path: &ObjectPath, + object_store: &ObjectStore, + ) -> lance_core::Result { + self.target + .resolve_latest_location(base_path, object_store) + .await + } + + async fn resolve_version_location( + &self, + base_path: &ObjectPath, + version: u64, + object_store: &dyn object_store::ObjectStore, + ) -> lance_core::Result { + let handler = if base_path == &self.source_base { + &self.source + } else { + &self.target + }; + handler + .resolve_version_location(base_path, version, object_store) + .await + } + + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + base_path: &ObjectPath, + object_store: &ObjectStore, + manifest_writer: ManifestWriter, + naming_scheme: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + self.target + .commit( + manifest, + indices, + base_path, + object_store, + manifest_writer, + naming_scheme, + transaction, + ) + .await + } + + async fn delete(&self, base_path: &ObjectPath) -> lance_core::Result<()> { + self.target.delete(base_path).await + } +} + /// Controls how new tables should be created #[derive(Clone, Debug, Default)] pub struct NewTableConfig { @@ -366,6 +440,27 @@ impl ListingDatabase { url.to_string() } + fn rooted_commit_handler_for_uri(uri: &str) -> Option> { + #[cfg(windows)] + { + let is_file = match url::Url::parse(uri) { + Ok(url) => url.scheme() == "file" || url.scheme().len() == 1, + Err(_) => true, + }; + return is_file.then(rooted_file_commit_handler); + } + + #[cfg(not(windows))] + { + let _ = uri; + None + } + } + + fn rooted_commit_handler(&self) -> Option> { + Self::rooted_commit_handler_for_uri(&self.uri) + } + async fn prepare_namespace_root( uri: &str, storage_options: &HashMap, @@ -763,7 +858,10 @@ impl ListingDatabase { if let Some(query_string) = &self.query_string { uri.push_str(&format!("?{}", query_string)); } - let commit_handler = commit_handler_from_url(&uri, &Some(object_store_params)).await?; + let commit_handler = match self.rooted_commit_handler() { + Some(handler) => handler, + None => commit_handler_from_url(&uri, &Some(object_store_params)).await?, + }; for name in names { let dir_name = format!("{}.{}", name, LANCE_EXTENSION); let full_path = self.base_path.clone().join(dir_name.clone()); @@ -896,6 +994,9 @@ impl ListingDatabase { write_params.mode = WriteMode::Overwrite; } + if write_params.commit_handler.is_none() { + write_params.commit_handler = self.rooted_commit_handler(); + } write_params.session = Some(self.session.clone()); write_params @@ -1142,9 +1243,11 @@ impl Database for ListingDatabase { }, ..Default::default() }; + let source_commit_handler = Self::rooted_commit_handler_for_uri(&request.source_uri); let read_params = ReadParams { store_options: Some(storage_params.clone()), session: Some(self.session.clone()), + commit_handler: source_commit_handler.clone(), ..Default::default() }; @@ -1170,7 +1273,7 @@ impl Database for ListingDatabase { let source_location = source_dataset .branch_location() .find_branch(ref_name.as_deref())?; - let (source_store, _) = ObjectStore::from_uri_and_params( + let (source_store, source_base) = ObjectStore::from_uri_and_params( self.session.store_registry(), &source_location.uri, &storage_params, @@ -1182,6 +1285,21 @@ impl Database for ListingDatabase { &storage_params, ) .await?; + let source_commit_handler = match source_commit_handler { + Some(handler) => handler, + None => { + commit_handler_from_url(&source_location.uri, &Some(storage_params.clone())).await? + } + }; + let target_commit_handler = match self.rooted_commit_handler() { + Some(handler) => handler, + None => commit_handler_from_url(&target_uri, &Some(storage_params.clone())).await?, + }; + let clone_commit_handler = Arc::new(CloneCommitHandler { + source_base, + source: source_commit_handler, + target: target_commit_handler, + }); let clone_op = Operation::Clone { is_shallow: true, ref_name, @@ -1194,6 +1312,7 @@ impl Database for ListingDatabase { .with_store_params(storage_params) .with_object_store(target_store) .with_source_store(source_store) + .with_commit_handler(clone_commit_handler) .with_session(self.session.clone()) .with_storage_format( source_dataset @@ -1276,6 +1395,9 @@ impl Database for ListingDatabase { } default_params }); + if read_params.commit_handler.is_none() { + read_params.commit_handler = self.rooted_commit_handler(); + } read_params.session(self.session.clone()); let native_table = Arc::new( @@ -1366,6 +1488,46 @@ mod tests { use std::path::PathBuf; use tempfile::tempdir; + #[derive(Debug)] + struct TargetOnlyCommitHandler; + + #[async_trait::async_trait] + impl CommitHandler for TargetOnlyCommitHandler { + async fn resolve_version_location( + &self, + _base_path: &ObjectPath, + _version: u64, + _object_store: &dyn object_store::ObjectStore, + ) -> lance_core::Result { + Err(lance_core::Error::invalid_input( + "target commit handler was used for the source manifest", + )) + } + + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + base_path: &ObjectPath, + object_store: &ObjectStore, + manifest_writer: ManifestWriter, + naming_scheme: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + lance_table::io::commit::ConditionalPutCommitHandler + .commit( + manifest, + indices, + base_path, + object_store, + manifest_writer, + naming_scheme, + transaction, + ) + .await + } + } + async fn setup_database() -> (tempfile::TempDir, ListingDatabase) { let tempdir = tempdir().unwrap(); let uri = tempdir.path().to_str().unwrap(); @@ -1481,6 +1643,75 @@ mod tests { assert!(matches!(reopened, Err(Error::TableNotFound { .. }))); } + #[cfg(windows)] + #[tokio::test] + async fn reopens_a_table_through_a_unc_authority() { + use std::path::{Component, Prefix}; + + let tempdir = tempdir().unwrap(); + let drive = match tempdir.path().components().next() { + Some(Component::Prefix(prefix)) => match prefix.kind() { + Prefix::Disk(drive) | Prefix::VerbatimDisk(drive) => drive, + other => panic!("expected a drive-backed temporary directory, got {other:?}"), + }, + other => panic!("expected an absolute Windows temporary directory, got {other:?}"), + }; + let drive_root = PathBuf::from(format!("{}:\\", drive as char)); + let relative = tempdir.path().strip_prefix(&drive_root).unwrap(); + let unc_path = PathBuf::from(format!(r"\\localhost\{}$", drive as char)).join(relative); + if !unc_path.try_exists().unwrap_or(false) { + eprintln!("skipping UNC lifecycle test because the localhost admin share is disabled"); + return; + } + let uri = url::Url::from_directory_path(&unc_path) + .unwrap() + .to_string(); + let request = ConnectRequest { + uri, + #[cfg(feature = "remote")] + client_config: Default::default(), + options: Default::default(), + namespace_client_properties: Default::default(), + manifest_enabled: false, + read_consistency_interval: None, + session: None, + }; + let db = ListingDatabase::connect_with_options(&request) + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let table = db + .create_table(CreateTableRequest { + name: "unc_reopen".to_string(), + namespace_path: vec![], + data: Box::new( + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]) + .unwrap(), + ), + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + }) + .await + .unwrap(); + drop(table); + + let reopened = db + .open_table(OpenTableRequest { + name: "unc_reopen".to_string(), + namespace_path: vec![], + index_cache_size: None, + lance_read_params: None, + location: None, + namespace_client: None, + managed_versioning: None, + }) + .await + .unwrap(); + assert_eq!(reopened.count_rows(None).await.unwrap(), 3); + } + #[tokio::test] async fn test_listing_database_root_ops_do_not_create_manifest() { let tempdir = tempdir().unwrap(); @@ -1764,6 +1995,114 @@ mod tests { assert!(!source_dir.path().join("target.lance").exists()); } + #[tokio::test] + async fn clone_routes_source_manifest_through_source_commit_handler() { + let source_dir = tempdir().unwrap(); + let target_dir = tempdir().unwrap(); + let source_request = ConnectRequest { + uri: source_dir.path().to_string_lossy().into_owned(), + #[cfg(feature = "remote")] + client_config: Default::default(), + options: Default::default(), + namespace_client_properties: Default::default(), + manifest_enabled: false, + read_consistency_interval: None, + session: None, + }; + let target_request = ConnectRequest { + uri: target_dir.path().to_string_lossy().into_owned(), + ..source_request.clone() + }; + let source_db = ListingDatabase::connect_with_options(&source_request) + .await + .unwrap(); + let target_db = ListingDatabase::connect_with_options(&target_request) + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + source_db + .create_table(CreateTableRequest { + name: "source".to_string(), + namespace_path: vec![], + data: Box::new( + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]) + .unwrap(), + ), + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + }) + .await + .unwrap(); + + let source_uri = source_db.table_uri("source").unwrap(); + let target_uri = target_db.table_uri("target").unwrap(); + let storage_params = ObjectStoreParams::default(); + let source_dataset = DatasetBuilder::from_uri(&source_uri) + .with_read_params(ReadParams { + session: Some(target_db.session.clone()), + ..Default::default() + }) + .load() + .await + .unwrap(); + let source_location = source_dataset.branch_location().find_branch(None).unwrap(); + let (source_store, source_base) = ObjectStore::from_uri_and_params( + target_db.session.store_registry(), + &source_location.uri, + &storage_params, + ) + .await + .unwrap(); + let (target_store, _) = ObjectStore::from_uri_and_params( + target_db.session.store_registry(), + &target_uri, + &storage_params, + ) + .await + .unwrap(); + let source_handler = commit_handler_from_url(&source_location.uri, &None) + .await + .unwrap(); + let clone_handler = Arc::new(CloneCommitHandler { + source_base, + source: source_handler, + target: Arc::new(TargetOnlyCommitHandler), + }); + let version = source_dataset.version().version; + let transaction = Transaction::new( + version, + Operation::Clone { + is_shallow: true, + ref_name: None, + ref_version: version, + ref_path: source_location.uri, + branch_name: None, + }, + None, + ); + + CommitBuilder::new(WriteDestination::Uri(&target_uri)) + .with_store_params(storage_params) + .with_object_store(target_store) + .with_source_store(source_store) + .with_commit_handler(clone_handler) + .with_session(target_db.session.clone()) + .with_storage_format( + source_dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + ) + .execute(transaction) + .await + .unwrap(); + + assert!(target_dir.path().join("target.lance").exists()); + } + #[tokio::test] async fn test_clone_table_with_storage_options() { let tempdir = tempdir().unwrap(); diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index 9f47c0b53..90ba560ca 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -5,8 +5,18 @@ use std::{fmt::Formatter, sync::Arc}; +#[cfg(any(windows, test))] +use futures::TryStreamExt; use futures::{StreamExt, TryFutureExt, stream::BoxStream}; use lance::io::WrappingObjectStore; +#[cfg(any(windows, test))] +use lance_table::{ + format::{IndexMetadata, Manifest, Transaction}, + io::commit::{ + CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, ManifestWriter, + RenameCommitHandler, + }, +}; use object_store::{ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, @@ -30,6 +40,59 @@ use url::Url; #[cfg(test)] pub mod io_tracking; +/// A local commit handler that resolves the latest manifest through the object store. +/// +/// Lance's native local shortcut reconstructs the selected manifest from a +/// filesystem path. On Windows that conversion drops the authority from a UNC +/// path. Listing through the already-rooted object store preserves the structural +/// server/share prefix while retaining the normal atomic-rename commit behavior. +#[derive(Debug)] +#[cfg(any(windows, test))] +struct RootedFileCommitHandler; + +#[cfg(any(windows, test))] +#[async_trait] +impl CommitHandler for RootedFileCommitHandler { + async fn resolve_latest_location( + &self, + base_path: &Path, + object_store: &lance::io::ObjectStore, + ) -> LanceResult { + self.list_manifest_locations(base_path, object_store, true) + .try_next() + .await? + .ok_or_else(|| LanceError::not_found(base_path.to_string())) + } + + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + base_path: &Path, + object_store: &lance::io::ObjectStore, + manifest_writer: ManifestWriter, + naming_scheme: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + RenameCommitHandler + .commit( + manifest, + indices, + base_path, + object_store, + manifest_writer, + naming_scheme, + transaction, + ) + .await + } +} + +#[cfg(any(windows, test))] +pub(crate) fn rooted_file_commit_handler() -> Arc { + Arc::new(RootedFileCommitHandler) +} + /// A file-store provider that anchors each request at its filesystem root. /// /// On Windows, an unprefixed [`LocalFileSystem`] cannot service UNC paths. Its @@ -269,12 +332,13 @@ impl ObjectStoreProvider for PrefixedFileStoreProvider { ) -> LanceResult { let (root, _) = Self::root_and_relative_path(&base_path)?; let raw_store: Arc = Arc::new(RootedLocalFileSystem::new(root)?); - // `file` keeps local planning and manifest behavior. With a custom - // wrapper, `file+uring` is also classified as local but (on Windows) - // routes readers and writers through the wrapped object store instead - // of the native fast path. + // Native local shortcuts are safe only when the rooted filesystem is + // the final store. An arbitrary wrapper can redirect I/O, so use + // Lance's non-cloud object-store route in that case. This keeps local + // scan planning without enabling native copy/delete or the io_uring + // scheduler, all of which would bypass the wrapper. let location = if params.object_store_wrapper.is_some() { - Url::parse("file+uring:///").expect("static URL must be valid") + Url::parse("memory:///").expect("static URL must be valid") } else { Url::parse("file:///").expect("static URL must be valid") }; @@ -344,6 +408,7 @@ pub(crate) fn new_prefixed_file_session() -> Arc { #[cfg(test)] mod prefixed_file_store_test { use super::*; + use object_store::memory::InMemory; use std::sync::atomic::{AtomicUsize, Ordering}; #[derive(Debug, Default)] @@ -362,6 +427,21 @@ mod prefixed_file_store_test { } } + #[derive(Debug, Default)] + struct MemoryRedirectWrapper { + store: Arc, + } + + impl WrappingObjectStore for MemoryRedirectWrapper { + fn wrap( + &self, + _store_prefix: &str, + _original: Arc, + ) -> Arc { + self.store.clone() + } + } + #[tokio::test] async fn anchors_new_and_existing_directories_at_a_filesystem_prefix() { let tempdir = tempfile::tempdir().unwrap(); @@ -437,9 +517,10 @@ mod prefixed_file_store_test { .await .unwrap(); assert_eq!(wrapper.calls.load(Ordering::Relaxed), 1); - assert_eq!(store.scheme(), "file+uring"); - assert!(store.is_local()); + assert_eq!(store.scheme(), "memory"); + assert!(!store.is_local()); assert!(!store.is_cloud()); + assert!(!store.prefers_lite_scheduler()); store .inner @@ -454,6 +535,90 @@ mod prefixed_file_store_test { assert_eq!(stats.written_bytes, 7); } + #[tokio::test] + async fn wrapped_store_cleanup_does_not_bypass_the_wrapper() { + let tempdir = tempfile::tempdir().unwrap(); + let table_path = tempdir.path().join("test.lance"); + std::fs::create_dir(&table_path).unwrap(); + std::fs::write(table_path.join("native-marker"), b"native").unwrap(); + + let table_url = Url::from_directory_path(&table_path).unwrap(); + let registry = Arc::new(ObjectStoreRegistry::default()); + registry.insert("file", Arc::new(PrefixedFileStoreProvider)); + let wrapper = Arc::new(MemoryRedirectWrapper::default()); + let params = ObjectStoreParams { + object_store_wrapper: Some(wrapper.clone()), + ..Default::default() + }; + let (store, base_path) = + lance::io::ObjectStore::from_uri_and_params(registry, table_url.as_str(), ¶ms) + .await + .unwrap(); + let wrapped_marker = base_path.clone().join("wrapped-marker"); + store + .inner + .put( + &wrapped_marker, + bytes::Bytes::from_static(b"wrapped").into(), + ) + .await + .unwrap(); + + store.remove_dir_all(base_path).await.unwrap(); + + assert!(table_path.join("native-marker").exists()); + assert!(matches!( + wrapper.store.head(&wrapped_marker).await, + Err(object_store::Error::NotFound { .. }) + )); + } + + #[tokio::test] + async fn rooted_commit_handler_uses_store_paths_for_latest_manifest() { + let tempdir = tempfile::tempdir().unwrap(); + let versions = tempdir.path().join("_versions"); + std::fs::create_dir(&versions).unwrap(); + std::fs::write(versions.join("2.manifest"), b"native").unwrap(); + + let base_path = Path::from_absolute_path(tempdir.path()).unwrap(); + let raw_store: Arc = Arc::new(InMemory::new()); + raw_store + .put( + &base_path.clone().join("_versions").join("1.manifest"), + bytes::Bytes::from_static(b"wrapped").into(), + ) + .await + .unwrap(); + let mut store = lance::io::ObjectStore::new( + raw_store.clone(), + Url::parse("file:///").unwrap(), + Some(4 * 1024), + None, + false, + false, + DEFAULT_LOCAL_IO_PARALLELISM, + 0, + None, + ); + store.inner = raw_store; + + let native = RenameCommitHandler + .resolve_latest_location(&base_path, &store) + .await + .unwrap(); + assert_eq!(native.version, 2); + + let rooted = rooted_file_commit_handler() + .resolve_latest_location(&base_path, &store) + .await + .unwrap(); + assert_eq!(rooted.version, 1); + assert_eq!( + rooted.path, + base_path.clone().join("_versions").join("1.manifest") + ); + } + #[tokio::test] async fn reuses_store_cache_across_database_paths() { let tempdir = tempfile::tempdir().unwrap();