diff --git a/docs/src/guide/object_store.md b/docs/src/guide/object_store.md index f901d2c2411..9ce4800b800 100644 --- a/docs/src/guide/object_store.md +++ b/docs/src/guide/object_store.md @@ -41,6 +41,35 @@ These options apply to all object stores. | `client_max_retries` | Number of times for the object store client to retry the request. Default, `3`. | | `client_retry_timeout` | Timeout for the object store client to retry the request in seconds. Default, `180`. | +## Per-Base Configuration + +A dataset can register additional base paths that store part of its data, and each +base may live in a different bucket, account, or storage provider. A storage option +key of the form `base_.` applies `` only to the base path with that +manifest id. Every base inherits the unscoped options; base-scoped entries add to or +override them for that base only. + +```python +import lance +ds = lance.dataset( + "az://account-a/path", + storage_options={ + # Shared defaults, used by the primary dataset and inherited by bases + "account_name": "account-a", + "account_key": "key-a", + # Overrides for the base path with id 1 + "base_1.account_name": "account-b", + "base_1.account_key": "key-b", + }, +) +``` + +Base ids are assigned when bases are registered (`initial_bases` ids are assigned +sequentially starting at 1, in order) and can be inspected through the manifest base +paths. Keys that do not match `base_.` exactly (e.g. `base_url`) are treated +as regular storage options. Exact per-base parameter maps (`base_store_params`, +keyed by base path URI) take precedence over base-scoped keys for that base. + ## S3 Configuration S3 (and S3-compatible stores) have additional configuration options that configure diff --git a/java/src/main/java/org/lance/ReadOptions.java b/java/src/main/java/org/lance/ReadOptions.java index a2e60e0e75b..5d6447539f0 100644 --- a/java/src/main/java/org/lance/ReadOptions.java +++ b/java/src/main/java/org/lance/ReadOptions.java @@ -189,6 +189,13 @@ public Builder setMetadataCacheSize(int metadataCacheSize) { * Set storage options. Extra options that make sense for a particular storage connection. This * is used to store connection parameters like credentials, endpoint, etc. * + *

For datasets with additional registered base paths, a key of the form {@code + * base_.} applies {@code } only to the base path with that manifest id, + * overriding the unscoped options that every base inherits. For example {@code + * base_1.account_key = abc} makes base 1 use {@code account_key = abc} while all other options + * are shared. Exact per-base bindings set via {@link #setBaseStoreParams(Map)} take precedence + * over base-scoped keys. + * * @param storageOptions the storage options * @return this builder */ diff --git a/java/src/main/java/org/lance/WriteParams.java b/java/src/main/java/org/lance/WriteParams.java index d8757006975..90d48feafb7 100644 --- a/java/src/main/java/org/lance/WriteParams.java +++ b/java/src/main/java/org/lance/WriteParams.java @@ -197,6 +197,17 @@ public Builder withDataStorageVersion(String dataStorageVersion) { return this; } + /** + * Set storage options for the write. + * + *

For writes involving additional registered base paths, a key of the form {@code + * base_.} applies {@code } only to the base path with that id, overriding the + * unscoped options that every base inherits. Exact per-base bindings set via {@link + * #withBaseStoreParams(Map)} take precedence over base-scoped keys. + * + * @param storageOptions the storage options + * @return this builder + */ public Builder withStorageOptions(Map storageOptions) { this.storageOptions = storageOptions; return this; diff --git a/python/python/lance/__init__.py b/python/python/lance/__init__.py index 4338950c649..0f929c1df32 100644 --- a/python/python/lance/__init__.py +++ b/python/python/lance/__init__.py @@ -166,6 +166,13 @@ def dataset( storage_options : optional, dict Extra options that make sense for a particular storage connection. This is used to store connection parameters like credentials, endpoint, etc. + + For datasets with additional registered base paths, a key of the form + ``base_.`` applies ```` only to the base path with that + manifest id, overriding the unscoped options that every base inherits. + For example ``{"account_key": "shared", "base_1.account_key": "abc"}`` + makes base 1 use ``account_key = abc`` while all other options are + shared. default_scan_options : optional, dict Default scan options that are used when scanning the dataset. This accepts the same arguments described in :py:meth:`lance.LanceDataset.scanner`. The @@ -203,9 +210,10 @@ def dataset( base_store_params : dict of str to dict, optional Runtime-only object store parameters keyed by base path URI. Each key is a base path URI (e.g., "s3://bucket/path") and each value is a dict - of storage options (credentials, endpoint, etc.) for that base. When a base - has no explicit entry here, the top-level ``storage_options`` is - used as a fallback. + of storage options (credentials, endpoint, etc.) for that base. These + take precedence over ``base_.`` entries in ``storage_options``. + When a base has no explicit entry here, the top-level + ``storage_options`` is used as a fallback. Notes ----- diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 98c09427f5e..76ba4eaa272 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7236,6 +7236,13 @@ def write_dataset( storage_options : optional, dict Extra options that make sense for a particular storage connection. This is used to store connection parameters like credentials, endpoint, etc. + + For writes involving additional base paths, a key of the form + ``base_.`` applies ```` only to the base path with that + id, overriding the unscoped options that every base inherits. For + example ``{"account_key": "shared", "base_1.account_key": "abc"}`` + makes base 1 use ``account_key = abc`` while all other options are + shared. data_storage_version: optional, str, default None The version of the data storage format to use. Newer versions are more efficient but require newer versions of lance to read. The default (None) @@ -7291,8 +7298,10 @@ def write_dataset( Runtime-only object store parameters keyed by base path URI. Each key is a base path URI (e.g., "s3://bucket/path") and each value is a dict of storage options (credentials, endpoint, etc.) for that base. These - are not persisted to the manifest. When a base has no explicit entry - here, the top-level ``storage_options`` is used as a fallback. + are not persisted to the manifest. These take precedence over + ``base_.`` entries in ``storage_options``. When a base has no + explicit entry here, the top-level ``storage_options`` is used as a + fallback. external_blob_mode: {"reference", "ingest"}, default "reference" How external blob URIs are handled on write. diff --git a/python/python/tests/test_multi_base.py b/python/python/tests/test_multi_base.py index 113f3224d14..3f5218f4cad 100644 --- a/python/python/tests/test_multi_base.py +++ b/python/python/tests/test_multi_base.py @@ -82,6 +82,40 @@ def test_multi_base_create_and_read(self): data.sort_values("id").reset_index(drop=True), ) + def test_base_scoped_storage_options(self): + """base_. storage options flow through write and read.""" + data = self.create_test_data(200) + + # Local stores ignore these options; this verifies base-scoped entries + # are resolved per base without breaking the write or read path. + storage_options = { + "shared_option": "shared", + "base_1.scoped_option": "base1-value", + } + + dataset = lance.write_dataset( + data, + self.primary_uri, + mode="create", + initial_bases=[DatasetBasePath(self.path1_uri, name="path1")], + target_bases=["path1"], + max_rows_per_file=100, + storage_options=storage_options, + ) + assert dataset.count_rows() == 200 + + # Data files land in the scoped base, not the primary path + assert list(Path(self.path1_uri).glob("**/*.lance")) + assert not list((Path(self.primary_uri) / "data").glob("*.lance")) + + # Reopen with the same flat options and read through the base store + dataset = lance.dataset(self.primary_uri, storage_options=storage_options) + result = dataset.to_table().to_pandas() + pd.testing.assert_frame_equal( + result.sort_values("id").reset_index(drop=True), + data.sort_values("id").reset_index(drop=True), + ) + def test_multi_base_append_mode(self): """Test appending data to a multi-base dataset.""" # Create initial dataset diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 8302801a24d..dafb5e5342f 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -3,6 +3,7 @@ //! Extend [object_store::ObjectStore] functionalities +use std::borrow::Cow; use std::collections::HashMap; use std::ops::Range; use std::pin::Pin; @@ -83,8 +84,10 @@ pub const DEFAULT_DOWNLOAD_RETRY_COUNT: usize = 3; pub use providers::{ObjectStoreProvider, ObjectStoreRegistry}; pub use storage_options::{ - EXPIRES_AT_MILLIS_KEY, LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, - StorageOptionsAccessor, StorageOptionsProvider, + BASE_SCOPED_OPTION_PREFIX, BaseScopedStorageOptionsProvider, EXPIRES_AT_MILLIS_KEY, + LanceNamespaceStorageOptionsProvider, REFRESH_OFFSET_MILLIS_KEY, StorageOptionsAccessor, + StorageOptionsProvider, has_base_scoped_options, parse_base_scoped_key, + resolve_base_scoped_options, }; #[async_trait] @@ -258,6 +261,27 @@ impl ObjectStoreParams { .as_ref() .and_then(|a| a.initial_storage_options()) } + + /// Resolve these params for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path; see + /// [`StorageOptionsAccessor::scoped_to_base`]. Returns the params unchanged + /// when the storage options contain no base-scoped entries. + pub fn scoped_to_base(&self, base_id: Option) -> Cow<'_, Self> { + let Some(accessor) = &self.storage_options_accessor else { + return Cow::Borrowed(self); + }; + let scoped = accessor.scoped_to_base(base_id); + if Arc::ptr_eq(&scoped, accessor) { + Cow::Borrowed(self) + } else { + Cow::Owned(Self { + storage_options_accessor: Some(scoped), + ..self.clone() + }) + } + } } // We implement hash for caching diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index c0b3a156502..775c98552a8 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -215,6 +215,12 @@ impl ObjectStoreRegistry { base_path: Url, params: &ObjectStoreParams, ) -> Result> { + // Base-scoped storage options (`base_.`) are directives for + // other registered base paths; resolve them away before building or + // caching a store for this location. Params already resolved for a + // base contain no scoped entries, so this is a no-op for them. + let params = params.scoped_to_base(None); + let params = params.as_ref(); let scheme = base_path.scheme(); let Some(provider) = self.get_provider(scheme) else { return Err(self.scheme_not_found_error(scheme)); @@ -394,6 +400,36 @@ mod tests { ); } + #[tokio::test] + async fn test_get_store_resolves_base_scoped_options() { + use crate::object_store::StorageOptionsAccessor; + + let registry = ObjectStoreRegistry::default(); + let url = Url::parse("memory://test").unwrap(); + + let with_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("shared".to_string(), "value".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ]), + ))), + ..Default::default() + }; + let without_scoped = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([("shared".to_string(), "value".to_string())]), + ))), + ..Default::default() + }; + + // Base-scoped entries are resolved away before the store is built and + // cached, so params with and without them yield the same cached store. + let store_scoped = registry.get_store(url.clone(), &with_scoped).await.unwrap(); + let store_plain = registry.get_store(url, &without_scoped).await.unwrap(); + assert!(Arc::ptr_eq(&store_scoped, &store_plain)); + } + #[test] fn test_calculate_object_store_scheme_not_found() { let registry = ObjectStoreRegistry::empty(); diff --git a/rust/lance-io/src/object_store/storage_options.rs b/rust/lance-io/src/object_store/storage_options.rs index 640e0a26b1e..fa5f4995afd 100644 --- a/rust/lance-io/src/object_store/storage_options.rs +++ b/rust/lance-io/src/object_store/storage_options.rs @@ -79,6 +79,15 @@ pub trait StorageOptionsProvider: Send + Sync + fmt::Debug { /// If [`EXPIRES_AT_MILLIS_KEY`] is not provided, the options are considered to never expire. async fn fetch_storage_options(&self) -> Result>>; + /// Fetch fresh storage options, bypassing caches along the chain. + /// + /// Providers that serve from an upstream cache (e.g. base-scoped wrappers + /// reading through a parent accessor) override this to force the upstream + /// to re-fetch. Defaults to [`Self::fetch_storage_options`]. + async fn force_fetch_storage_options(&self) -> Result>> { + self.fetch_storage_options().await + } + /// Return a human-readable unique identifier for this provider instance /// /// This is used for equality comparison and hashing in the object store registry. @@ -155,6 +164,104 @@ impl StorageOptionsProvider for LanceNamespaceStorageOptionsProvider { } } +/// Prefix marking a storage option as scoped to a single registered base path. +/// +/// A key of the form `base_.` applies `` only to the base path +/// with manifest id ``, overriding the shared (unscoped) options for that +/// base. For example `base_1.account_key = abc` gives the base with id 1 the +/// option `account_key = abc` while it inherits every unscoped option. +pub const BASE_SCOPED_OPTION_PREFIX: &str = "base_"; + +/// Parse a base-scoped storage option key of the form `base_.`. +/// +/// Returns `Some((base_id, key))` only for keys that match the convention +/// exactly: the `base_` prefix, an all-digit u32 base id, a `.` separator, and +/// a non-empty remainder. Any other key (e.g. `base_url`, `base_x.key`, +/// `base_1.`) is not base-scoped. +pub fn parse_base_scoped_key(key: &str) -> Option<(u32, &str)> { + let rest = key.strip_prefix(BASE_SCOPED_OPTION_PREFIX)?; + let (id_str, scoped_key) = rest.split_once('.')?; + if scoped_key.is_empty() || id_str.is_empty() || !id_str.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + let id = id_str.parse::().ok()?; + Some((id, scoped_key)) +} + +/// Returns true if any key in `options` is base-scoped (`base_.`). +pub fn has_base_scoped_options(options: &HashMap) -> bool { + options + .keys() + .any(|key| parse_base_scoped_key(key).is_some()) +} + +/// Resolve the effective storage options for one base path scope. +/// +/// All base-scoped keys are removed from the result. When `base_id` is +/// `Some(id)`, entries scoped to that id are overlaid on the unscoped options, +/// adding or overriding keys. `None` resolves the default scope (the primary +/// dataset base), which simply drops every base-scoped entry. +pub fn resolve_base_scoped_options( + options: &HashMap, + base_id: Option, +) -> HashMap { + let mut resolved = HashMap::with_capacity(options.len()); + let mut overrides = Vec::new(); + for (key, value) in options { + match parse_base_scoped_key(key) { + Some((id, scoped_key)) => { + if Some(id) == base_id { + overrides.push((scoped_key.to_string(), value.clone())); + } + } + None => { + resolved.insert(key.clone(), value.clone()); + } + } + } + resolved.extend(overrides); + resolved +} + +/// A [`StorageOptionsProvider`] that resolves another accessor's options for a +/// single base path scope. +/// +/// Fetching through this provider first refreshes the parent accessor when its +/// options have expired, then resolves the refreshed options for the scope. As +/// a result, dynamically vended per-base credentials (e.g. a namespace server +/// returning `base_.` entries in one flat map) stay fresh per base. +#[derive(Debug)] +pub struct BaseScopedStorageOptionsProvider { + inner: Arc, + base_id: Option, +} + +impl BaseScopedStorageOptionsProvider { + pub fn new(inner: Arc, base_id: Option) -> Self { + Self { inner, base_id } + } +} + +#[async_trait] +impl StorageOptionsProvider for BaseScopedStorageOptionsProvider { + async fn fetch_storage_options(&self) -> Result>> { + let options = self.inner.get_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + async fn force_fetch_storage_options(&self) -> Result>> { + let options = self.inner.refresh_storage_options().await?; + Ok(Some(resolve_base_scoped_options(&options.0, self.base_id))) + } + + fn provider_id(&self) -> String { + match self.base_id { + Some(id) => format!("base-scoped[base_id={}]({})", id, self.inner.accessor_id()), + None => format!("base-scoped[default]({})", self.inner.accessor_id()), + } + } +} + /// Unified access to storage options with automatic caching and refresh /// /// This struct bundles static storage options with an optional dynamic provider, @@ -184,6 +291,11 @@ pub struct StorageOptionsAccessor { /// Duration before expiry to trigger refresh refresh_offset: Duration, + + /// True when this accessor was produced by [`Self::scoped_to_base`]; its + /// options are already resolved for one base path scope, so scoping again + /// is a no-op. + scope_resolved: bool, } impl fmt::Debug for StorageOptionsAccessor { @@ -212,14 +324,33 @@ impl StorageOptionsAccessor { .unwrap_or(Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS)) } + /// Effective expiration of a raw options map: the minimum of the unscoped + /// `expires_at_millis` and every `base_.expires_at_millis` entry. + /// + /// A flat map may vend per-base credentials that expire before the shared + /// ones. Refreshing when the earliest credential is due keeps base-scoped + /// accessors (which refresh through this accessor) from re-resolving stale + /// per-base credentials out of a still-"valid" cache. + fn effective_expires_at_millis(options: &HashMap) -> Option { + options + .iter() + .filter(|(key, _)| { + key.as_str() == EXPIRES_AT_MILLIS_KEY + || matches!( + parse_base_scoped_key(key), + Some((_, scoped_key)) if scoped_key == EXPIRES_AT_MILLIS_KEY + ) + }) + .filter_map(|(_, value)| value.parse::().ok()) + .min() + } + /// Create an accessor with only static options (no refresh capability) /// /// The returned accessor will always return the provided options. /// This is useful when credentials don't expire or are managed externally. pub fn with_static_options(options: HashMap) -> Self { - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); let refresh_offset = Self::extract_refresh_offset(&options); Self { @@ -230,6 +361,7 @@ impl StorageOptionsAccessor { expires_at_millis, }))), refresh_offset, + scope_resolved: false, } } @@ -247,6 +379,7 @@ impl StorageOptionsAccessor { provider: Some(provider), cache: Arc::new(RwLock::new(None)), refresh_offset: Duration::from_millis(DEFAULT_REFRESH_OFFSET_MILLIS), + scope_resolved: false, } } @@ -263,9 +396,7 @@ impl StorageOptionsAccessor { initial_options: HashMap, provider: Arc, ) -> Self { - let expires_at_millis = initial_options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&initial_options); let refresh_offset = Self::extract_refresh_offset(&initial_options); Self { @@ -276,6 +407,7 @@ impl StorageOptionsAccessor { expires_at_millis, }))), refresh_offset, + scope_resolved: false, } } @@ -306,8 +438,9 @@ impl StorageOptionsAccessor { /// Fetch fresh options from the provider and update the cache. /// /// This bypasses the cache for callers that need to validate provider-vended - /// credentials even when initial metadata has no expiration. - #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] + /// credentials even when initial metadata has no expiration. The force + /// propagates through provider chains (e.g. base-scoped wrappers), so the + /// origin provider is re-queried even when intermediate caches are valid. pub(crate) async fn refresh_storage_options(&self) -> Result { let Some(provider) = &self.provider else { return self.get_storage_options().await; @@ -318,7 +451,7 @@ impl StorageOptionsAccessor { provider.provider_id() ); - let storage_options_map = provider.fetch_storage_options().await.map_err(|e| { + let storage_options_map = provider.force_fetch_storage_options().await.map_err(|e| { Error::io_source(Box::new(std::io::Error::other(format!( "Failed to fetch storage options: {}", e @@ -336,9 +469,7 @@ impl StorageOptionsAccessor { return Ok(super::StorageOptions(HashMap::new())); }; - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); let mut cache = self.cache.write().await; *cache = Some(CachedStorageOptions { @@ -409,9 +540,7 @@ impl StorageOptionsAccessor { return Ok(Some(super::StorageOptions(HashMap::new()))); }; - let expires_at_millis = options - .get(EXPIRES_AT_MILLIS_KEY) - .and_then(|s| s.parse::().ok()); + let expires_at_millis = Self::effective_expires_at_millis(&options); if let Some(expires_at) = expires_at_millis { let now_ms = SystemTime::now() @@ -493,6 +622,50 @@ impl StorageOptionsAccessor { } } + /// Resolve this accessor for a single base path scope. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to one registered base path. The returned accessor resolves + /// options for `base_id`: entries scoped to that base overlay the unscoped + /// defaults, and all other scoped entries are dropped. `None` resolves the + /// default scope used for the primary dataset base. + /// + /// A static accessor whose options contain no base-scoped entries is + /// returned unchanged, preserving accessor identity (and thus object store + /// registry cache keys). A provider-backed accessor is always wrapped + /// through [`BaseScopedStorageOptionsProvider`] — fetched options may vend + /// base-scoped entries at any refresh, even when the initial options carry + /// none — so refreshed options are re-resolved for the scope on every + /// fetch. Accessors already produced by this method are returned unchanged. + pub fn scoped_to_base(self: &Arc, base_id: Option) -> Arc { + if self.scope_resolved { + return self.clone(); + } + if self.has_provider() { + let provider = Arc::new(BaseScopedStorageOptionsProvider::new(self.clone(), base_id)); + let mut scoped = match self.initial_storage_options() { + Some(initial) => Self::with_initial_and_provider( + resolve_base_scoped_options(initial, base_id), + provider, + ), + None => Self::with_provider(provider), + }; + scoped.scope_resolved = true; + Arc::new(scoped) + } else { + match self.initial_storage_options() { + Some(initial) if has_base_scoped_options(initial) => { + let mut scoped = + Self::with_static_options(resolve_base_scoped_options(initial, base_id)); + scoped.scope_resolved = true; + Arc::new(scoped) + } + // Static options never change, so there is nothing to scope. + _ => self.clone(), + } + } + } + /// Check if this accessor has a dynamic provider pub fn has_provider(&self) -> bool { self.provider.is_some() @@ -748,4 +921,322 @@ mod tests { accessor.get_storage_options().await.unwrap(); assert_eq!(mock_provider.get_call_count().await, 1); } + + #[test] + fn test_parse_base_scoped_key() { + assert_eq!( + parse_base_scoped_key("base_1.account_key"), + Some((1, "account_key")) + ); + assert_eq!( + parse_base_scoped_key("base_12.headers.x-ms-version"), + Some((12, "headers.x-ms-version")) + ); + assert_eq!(parse_base_scoped_key("base_0.region"), Some((0, "region"))); + + // Not base-scoped keys + assert_eq!(parse_base_scoped_key("account_key"), None); + assert_eq!(parse_base_scoped_key("base_url"), None); + assert_eq!(parse_base_scoped_key("base_hot.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1x.account_key"), None); + assert_eq!(parse_base_scoped_key("base_+1.account_key"), None); + assert_eq!(parse_base_scoped_key("base_.account_key"), None); + assert_eq!(parse_base_scoped_key("base_1."), None); + assert_eq!(parse_base_scoped_key("base_1"), None); + // Overflows u32 + assert_eq!(parse_base_scoped_key("base_4294967296.key"), None); + } + + #[test] + fn test_resolve_base_scoped_options() { + let options = HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ("base_2.account_key".to_string(), "base2-key".to_string()), + ("base_2.endpoint".to_string(), "http://b2".to_string()), + ]); + assert!(has_base_scoped_options(&options)); + + let base1 = resolve_base_scoped_options(&options, Some(1)); + assert_eq!( + base1, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base1-key".to_string()), + ]) + ); + + let base2 = resolve_base_scoped_options(&options, Some(2)); + assert_eq!( + base2, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "base2-key".to_string()), + ("endpoint".to_string(), "http://b2".to_string()), + ]) + ); + + // A base without scoped entries inherits only the unscoped options + let base3 = resolve_base_scoped_options(&options, Some(3)); + assert_eq!( + base3, + HashMap::from([ + ("region".to_string(), "us-east-1".to_string()), + ("account_key".to_string(), "shared-key".to_string()), + ]) + ); + + // The default scope drops every scoped entry + let default = resolve_base_scoped_options(&options, None); + assert_eq!(default, base3); + + assert!(!has_base_scoped_options(&HashMap::from([( + "account_key".to_string(), + "shared-key".to_string() + )]))); + } + + #[tokio::test] + async fn test_scoped_to_base_identity_and_idempotency() { + // Static accessors without scoped keys are returned unchanged. + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [("account_key".to_string(), "shared-key".to_string())], + ))); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(Some(1)), &accessor)); + assert!(Arc::ptr_eq(&accessor.scoped_to_base(None), &accessor)); + + // Scoping an already-scoped accessor is a no-op (the registry choke + // point re-applies the default scope to every params it sees). + let scoped = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq(&scoped.scoped_to_base(None), &scoped)); + + let provider_scoped = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + MockStorageOptionsProvider::new(None), + ))) + .scoped_to_base(Some(1)); + assert!(Arc::ptr_eq( + &provider_scoped.scoped_to_base(None), + &provider_scoped + )); + } + + #[tokio::test] + async fn test_scoped_to_base_provider_only_resolves_vended_options() { + MockClock::set_system_time(Duration::from_secs(100_000)); + + // No initial options: scoped entries arrive only through the provider. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let parent = Arc::new(StorageOptionsAccessor::with_provider(provider.clone())); + + let base1 = parent.scoped_to_base(Some(1)); + assert!(!Arc::ptr_eq(&base1, &parent)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + let default = parent.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert!(!result.0.contains_key("base_1.account_key")); + + // Both scopes were served from one parent fetch. + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_earlier_base_expiry_refreshes_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + // Base 1 credentials expire before the shared ones; the parent must + // refresh when the earliest credential is due, or the scoped accessor + // would keep re-resolving stale base-1 credentials from a still- + // "valid" parent cache. + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ( + format!("base_1.{}", EXPIRES_AT_MILLIS_KEY), + (now_ms + 120_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Past the base-1 expiry but before the shared expiry: the parent's + // effective expiry is the earlier one, so the refresh chain fetches + // fresh credentials instead of re-serving BASE1_0. + MockClock::set_system_time(Duration::from_secs(100_000 + 121)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_to_base_static() { + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [ + ("account_key".to_string(), "shared-key".to_string()), + ("base_1.account_key".to_string(), "base1-key".to_string()), + ], + ))); + + let base1 = accessor.scoped_to_base(Some(1)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "base1-key".to_string())]) + ); + assert!(!base1.has_provider()); + + let default = accessor.scoped_to_base(None); + let result = default.get_storage_options().await.unwrap(); + assert_eq!( + result.0, + HashMap::from([("account_key".to_string(), "shared-key".to_string())]) + ); + + // Scoped accessor ids are stable across derivations and distinct per scope + assert_eq!( + accessor.scoped_to_base(Some(1)).accessor_id(), + base1.accessor_id() + ); + assert_ne!(base1.accessor_id(), default.accessor_id()); + assert_ne!(base1.accessor_id(), accessor.accessor_id()); + } + + #[derive(Debug)] + struct MockBaseScopedVendingProvider { + call_count: Arc>, + expires_in_millis: u64, + } + + #[async_trait] + impl StorageOptionsProvider for MockBaseScopedVendingProvider { + async fn fetch_storage_options(&self) -> Result>> { + let count = { + let mut c = self.call_count.write().await; + *c += 1; + *c + }; + let now_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + Ok(Some(HashMap::from([ + ("account_key".to_string(), format!("SHARED_{}", count)), + ("base_1.account_key".to_string(), format!("BASE1_{}", count)), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + self.expires_in_millis).to_string(), + ), + ]))) + } + + fn provider_id(&self) -> String { + "MockBaseScopedVendingProvider".to_string() + } + } + + #[tokio::test] + async fn test_scoped_to_base_refreshes_through_parent() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + + let base1 = parent.scoped_to_base(Some(1)); + let default = parent.scoped_to_base(None); + assert!(base1.has_provider()); + + // Initial options are resolved per scope without fetching + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_0"); + assert!(!result.0.contains_key("base_1.account_key")); + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_0"); + assert_eq!(*provider.call_count.read().await, 0); + + // Expire the vended options; the scoped accessor refreshes through the + // parent and re-resolves the refreshed options for its scope. + MockClock::set_system_time(Duration::from_secs(100_000 + 601)); + let result = base1.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + + // The parent refresh is shared: other scopes see it without refetching + let result = default.get_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "SHARED_1"); + assert_eq!(*provider.call_count.read().await, 1); + } + + #[tokio::test] + async fn test_scoped_forced_refresh_reaches_origin_provider() { + MockClock::set_system_time(Duration::from_secs(100_000)); + let now_ms = MockClock::system_time().as_millis() as u64; + + let provider = Arc::new(MockBaseScopedVendingProvider { + call_count: Arc::new(RwLock::new(0)), + expires_in_millis: 600_000, + }); + let initial = HashMap::from([ + ("account_key".to_string(), "SHARED_0".to_string()), + ("base_1.account_key".to_string(), "BASE1_0".to_string()), + ( + EXPIRES_AT_MILLIS_KEY.to_string(), + (now_ms + 600_000).to_string(), + ), + ]); + let parent = Arc::new(StorageOptionsAccessor::with_initial_and_provider( + initial, + provider.clone(), + )); + let base1 = parent.scoped_to_base(Some(1)); + + // A forced refresh must reach the origin provider even though both the + // scoped and the parent caches are still valid. + let result = base1.refresh_storage_options().await.unwrap(); + assert_eq!(result.0.get("account_key").unwrap(), "BASE1_1"); + assert_eq!(*provider.call_count.read().await, 1); + } } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 9f6a683db09..6a0440f6804 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -806,6 +806,11 @@ impl Dataset { let metadata_cache = Arc::new(session.metadata_cache.for_dataset(&uri)); let index_cache = Arc::new(session.index_cache.for_dataset(&uri)); let fragment_bitmap = Arc::new(manifest.fragments.iter().map(|f| f.id as u32).collect()); + write::log_unregistered_base_scoped_options( + store_params.as_ref(), + &manifest.base_paths, + log::Level::Debug, + ); Ok(Self { object_store, base: base_path, @@ -1886,21 +1891,26 @@ impl Dataset { store_params } - fn store_params_for_base( + pub(crate) fn store_params_for_base( &self, base_path: Option<&lance_table::format::BasePath>, ) -> ObjectStoreParams { // Base-specific bindings are exact ObjectStoreParams keyed by - // `BasePath.path`. If a base has no explicit binding then reads fall back - // to the dataset-level default store params. - base_path - .and_then(|base_path| { - self.base_store_params - .as_ref() - .and_then(|params| params.get(&base_path.path)) - }) - .cloned() - .unwrap_or_else(|| self.store_params.as_deref().cloned().unwrap_or_default()) + // `BasePath.path` and are used as-is. Otherwise the dataset-level + // default params are resolved for the base scope: `base_.` + // storage options overlay the shared defaults for that base. + if let Some(params) = base_path.and_then(|base_path| { + self.base_store_params + .as_ref() + .and_then(|params| params.get(&base_path.path)) + }) { + return params.clone(); + } + let default_params = self.store_params.as_deref().cloned().unwrap_or_default(); + match default_params.scoped_to_base(base_path.map(|base_path| base_path.id)) { + Cow::Owned(scoped_params) => scoped_params, + Cow::Borrowed(_) => default_params, + } } /// Returns the initial storage options used when opening this dataset, if any. diff --git a/rust/lance/src/dataset/builder.rs b/rust/lance/src/dataset/builder.rs index 3f3a166b9c1..aecbf92f0d8 100644 --- a/rust/lance/src/dataset/builder.rs +++ b/rust/lance/src/dataset/builder.rs @@ -305,6 +305,14 @@ impl DatasetBuilder { /// - [Azure options](https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html#variants) /// - [S3 options](https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html#variants) /// - [Google options](https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html#variants) + /// + /// For datasets with additional registered base paths, a key of the form + /// `base_.` applies `` only to the base path with that + /// manifest id, overriding the unscoped options that every base inherits. + /// For example `base_1.account_key = abc` makes base 1 use + /// `account_key = abc` while all other options are shared. Exact per-base + /// bindings set via [`Self::with_base_store_params`] take precedence over + /// base-scoped keys. pub fn with_storage_options(mut self, storage_options: HashMap) -> Self { // Merge with existing options if accessor exists, otherwise create new static accessor if let Some(existing) = self.options.storage_options_accessor.take() { @@ -459,8 +467,9 @@ impl DatasetBuilder { /// /// These params are not persisted in the manifest. They are used as-is /// whenever the dataset resolves an object store for the given - /// `BasePath.path`. Dataset-level store params remain the fallback for bases - /// without an explicit binding. + /// `BasePath.path`, taking precedence over `base_.` storage + /// options (see [`Self::with_storage_options`]). Dataset-level store params + /// remain the fallback for bases without an explicit binding. pub fn with_base_store_params( mut self, base_path: impl AsRef, diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 8599db9e945..526c9aaeb64 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -241,6 +241,62 @@ async fn test_with_object_store_wrappers_wraps_base_store_params() { assert!(request_tracker.incremental_stats().read_iops > 0); } +#[tokio::test] +async fn test_store_params_for_base_resolves_base_scoped_options() { + let test_dir = TempStdDir::default(); + create_file(&test_dir, WriteMode::Create, LanceFileVersion::Stable).await; + let uri = test_dir.to_str().unwrap(); + let dataset = Arc::new(Dataset::open(uri).await.unwrap()); + + let base_dir = tempfile::tempdir().unwrap(); + let base_uri = file_object_store_uri(base_dir.path()); + let base = BasePath::new(1, base_uri, Some("base".to_string()), true); + dataset.add_bases(vec![base.clone()], None).await.unwrap(); + + // Reopen with a single flat storage options map carrying base-scoped + // entries (`base_.`) next to shared defaults. + let dataset = DatasetBuilder::from_uri(uri) + .with_storage_options(HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ( + "base_1.scoped_option".to_string(), + "base1-value".to_string(), + ), + ])) + .load() + .await + .unwrap(); + + // The registered base resolves the scoped entry on top of shared defaults. + let base_path = dataset.manifest.base_paths.get(&1).unwrap().clone(); + let params = dataset.store_params_for_base(Some(&base_path)); + assert_eq!( + params.storage_options().unwrap(), + &HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ("scoped_option".to_string(), "base1-value".to_string()), + ]) + ); + + // The default scope keeps only the shared defaults. + let params = dataset.store_params_for_base(None); + assert_eq!( + params.storage_options().unwrap(), + &HashMap::from([("shared_option".to_string(), "shared".to_string())]) + ); + + // The base store resolved from scoped options is usable end to end. + let base_store = dataset.object_store(Some(1)).await.unwrap(); + let probe = base + .extract_path(dataset.session().store_registry()) + .unwrap() + .join("data") + .join("probe.lance"); + base_store.put(&probe, b"hello").await.unwrap(); + let read = base_store.inner.get_range(&probe, 0..5).await.unwrap(); + assert_eq!(read.as_ref(), b"hello"); +} + #[tokio::test] async fn test_with_object_store_wrappers_wraps_refs_store() { let test_dir = tempfile::tempdir().unwrap(); diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 9633b26698f..167b2f70b89 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -25,12 +25,15 @@ use lance_file::previous::writer::{ }; use lance_file::version::LanceFileVersion; use lance_file::writer::{self as current_writer, FileWriterOptions}; -use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry}; +use lance_io::object_store::{ + ObjectStore, ObjectStoreParams, ObjectStoreRegistry, parse_base_scoped_key, +}; use lance_table::format::{BasePath, DataFile, Fragment}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; use object_store::path::Path; -use std::collections::{HashMap, HashSet}; +use std::borrow::Cow; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::num::NonZero; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -287,8 +290,16 @@ pub struct WriteParams { /// Write mode pub mode: WriteMode, + /// Default object store params for the write. + /// + /// Storage options may carry base-scoped entries (`base_.`) that + /// apply only to the registered base path with that id, overriding the + /// unscoped options that every base inherits. pub store_params: Option, + /// Exact object store params per base path URI, taking precedence over + /// `base_.` storage options in [`Self::store_params`]. See + /// [`Self::with_base_store_params`]. pub base_store_params: Option>, pub progress: Arc, @@ -445,8 +456,10 @@ impl WriteParams { /// Set exact runtime object store params for a registered base path. /// - /// These params are used as-is for that base. The write-level default - /// `store_params` remain the fallback for bases without an explicit binding. + /// These params are used as-is for that base, taking precedence over + /// `base_.` storage options in `store_params`. The write-level + /// default `store_params` remain the fallback for bases without an + /// explicit binding. pub fn with_base_store_params( mut self, base_path: impl AsRef, @@ -798,6 +811,11 @@ pub async fn validate_and_resolve_target_bases( all_bases.insert(base_path.id, base_path.clone()); } } + log_unregistered_base_scoped_options( + params.store_params.as_ref(), + &all_bases, + log::Level::Warn, + ); // Step 3: Resolve target_base_names_or_paths to IDs let target_base_ids = if let Some(ref names_or_paths) = params.target_base_names_or_paths { @@ -843,7 +861,7 @@ pub async fn validate_and_resolve_target_bases( )) })?; - let store_params = write_store_params_for_base(params, &base_path.path); + let store_params = write_store_params_for_base(params, base_path); let (target_object_store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -886,22 +904,53 @@ fn append_external_base_candidate( } } -fn write_store_params_for_base(params: &WriteParams, base_path: &str) -> ObjectStoreParams { - params - .base_store_params - .as_ref() - .and_then(|base_store_params| base_store_params.get(base_path)) - .cloned() - .unwrap_or_else(|| params.store_params.clone().unwrap_or_default()) +/// Log base-scoped storage options (`base_.`) whose id does not +/// match any registered base path. Unregistered entries are ignored during +/// resolution. The open path logs at debug, since options may legitimately be +/// vended for bases the loaded version does not register; the write path logs +/// at warn, since ids are already assigned there and an unmatched id is much +/// more likely a mistake. +pub(crate) fn log_unregistered_base_scoped_options( + store_params: Option<&ObjectStoreParams>, + base_paths: &HashMap, + level: log::Level, +) { + if !log::log_enabled!(level) { + return; + } + let Some(options) = store_params.and_then(|params| params.storage_options()) else { + return; + }; + let unregistered = options + .keys() + .filter_map(|key| parse_base_scoped_key(key).map(|(id, _)| id)) + .filter(|id| !base_paths.contains_key(id)) + .collect::>(); + if !unregistered.is_empty() { + log::log!( + level, + "Ignoring base-scoped storage options for unregistered base path ids: {:?}", + unregistered + ); + } } -fn dataset_store_params_for_base(dataset: &Dataset, base_path: &str) -> ObjectStoreParams { - dataset +fn write_store_params_for_base(params: &WriteParams, base_path: &BasePath) -> ObjectStoreParams { + // Exact per-URI bindings are used as-is. Otherwise the write-level default + // params are resolved for the base scope: `base_.` storage + // options overlay the shared defaults for that base. + if let Some(store_params) = params .base_store_params .as_ref() - .and_then(|base_store_params| base_store_params.get(base_path)) - .cloned() - .unwrap_or_else(|| dataset.store_params.as_deref().cloned().unwrap_or_default()) + .and_then(|base_store_params| base_store_params.get(&base_path.path)) + { + return store_params.clone(); + } + let default_params = params.store_params.clone().unwrap_or_default(); + match default_params.scoped_to_base(Some(base_path.id)) { + Cow::Owned(scoped_params) => scoped_params, + Cow::Borrowed(_) => default_params, + } } async fn append_external_initial_bases( @@ -913,7 +962,7 @@ async fn append_external_initial_bases( ) -> Result<()> { if let Some(initial_bases) = initial_bases { for base_path in initial_bases { - let store_params = write_store_params_for_base(params, &base_path.path); + let store_params = write_store_params_for_base(params, base_path); let (store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -946,7 +995,7 @@ async fn build_external_base_resolver( if let Some(dataset) = dataset { for base_path in dataset.manifest.base_paths.values() { - let store_params = dataset_store_params_for_base(dataset, &base_path.path); + let store_params = dataset.store_params_for_base(Some(base_path)); let (store, extracted_path) = ObjectStore::from_uri_and_params( store_registry.clone(), &base_path.path, @@ -2103,7 +2152,28 @@ mod tests { .with_base_store_params("az://container/path-a", azure_store_params("account-a")) .with_base_store_params("az://container/path-b", azure_store_params("account-b")); - let existing_base_paths = HashMap::from([ + let existing_base_paths = azure_base_paths_a_b(); + + let target_bases = + validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) + .await + .unwrap() + .unwrap(); + + assert_eq!(target_bases.len(), 2); + assert_eq!( + target_bases[0].object_store.store_prefix, + "az$container@account-a" + ); + assert_eq!( + target_bases[1].object_store.store_prefix, + "az$container@account-b" + ); + } + + #[cfg(feature = "azure")] + fn azure_base_paths_a_b() -> HashMap { + HashMap::from([ ( 1, BasePath::new( @@ -2122,7 +2192,32 @@ mod tests { false, ), ), - ]); + ]) + } + + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_validate_and_resolve_target_bases_uses_base_scoped_storage_options() { + // A single flat storage options map carries per-base credentials via + // the `base_.` convention; unscoped keys are shared defaults. + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_name".to_string(), "account-shared".to_string()), + ("account_key".to_string(), "dGVzdA==".to_string()), + ("base_1.account_name".to_string(), "account-a".to_string()), + ("base_2.account_name".to_string(), "account-b".to_string()), + ]), + ))), + ..Default::default() + }; + let mut params = WriteParams { + store_params: Some(store_params), + ..Default::default() + } + .with_target_bases(vec![1, 2]); + + let existing_base_paths = azure_base_paths_a_b(); let target_bases = validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) @@ -2141,6 +2236,43 @@ mod tests { ); } + #[cfg(feature = "azure")] + #[tokio::test] + async fn test_base_store_params_take_precedence_over_base_scoped_options() { + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("account_key".to_string(), "dGVzdA==".to_string()), + ( + "base_1.account_name".to_string(), + "account-scoped".to_string(), + ), + ]), + ))), + ..Default::default() + }; + let mut params = WriteParams { + store_params: Some(store_params), + ..Default::default() + } + .with_target_bases(vec![1]) + .with_base_store_params("az://container/path-a", azure_store_params("account-exact")); + + let existing_base_paths = azure_base_paths_a_b(); + + let target_bases = + validate_and_resolve_target_bases(&mut params, Some(&existing_base_paths)) + .await + .unwrap() + .unwrap(); + + assert_eq!(target_bases.len(), 1); + assert_eq!( + target_bases[0].object_store.store_prefix, + "az$container@account-exact" + ); + } + #[tokio::test] async fn test_explicit_data_file_bases_writer_generator() { use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema}; @@ -2538,6 +2670,80 @@ mod tests { ); } + #[tokio::test] + async fn test_multi_base_write_read_with_base_scoped_storage_options() { + use crate::dataset::builder::DatasetBuilder; + use lance_core::utils::tempfile::TempStrDir; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + + let primary_dir = TempStrDir::default(); + let base1_dir = TempStrDir::default(); + + // Local stores ignore these options; this verifies base-scoped entries + // flow through the full write/read path without breaking anything. + let scoped_options = HashMap::from([ + ("shared_option".to_string(), "shared".to_string()), + ( + "base_1.scoped_option".to_string(), + "base1-value".to_string(), + ), + ]); + let store_params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + scoped_options.clone(), + ))), + ..Default::default() + }; + + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + let dataset = Dataset::write( + data_gen.batch(5), + primary_dir.as_str(), + Some(WriteParams { + mode: WriteMode::Create, + store_params: Some(store_params), + initial_bases: Some(vec![BasePath { + id: 1, + name: Some("base1".to_string()), + path: base1_dir.as_str().to_string(), + is_dataset_root: true, + }]), + target_bases: Some(vec![1]), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 5); + for fragment in dataset.get_fragments() { + assert!( + fragment + .metadata + .files + .iter() + .all(|file| file.base_id == Some(1)) + ); + } + + // Reopen with the same flat options and scan through the base store. + let dataset = DatasetBuilder::from_uri(primary_dir.as_str()) + .with_storage_options(scoped_options) + .load() + .await + .unwrap(); + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(num_rows, 5); + } + #[tokio::test] async fn test_multi_base_overwrite() { use lance_testing::datagen::{BatchGenerator, IncrementingInt32};