fix(rust): preserve AWS store provider contracts

rust/lancedb/src/database/listing.rs

@@ -24,7 +24,7 @@ use crate::database::ReadConsistency;
 use crate::database::namespace::LanceNamespaceDatabase;
 use crate::error::{CreateDirSnafu, Error, Result};
 use crate::io::object_store::{
-    MirroringObjectStoreWrapper, install_atomic_aws_provider,
+    MirroringObjectStoreWrapper, install_atomic_aws_provider, is_aws_credential_option,
     object_store_params_from_storage_options, set_storage_options,
 };
 use crate::table::NativeTable;
@@ -739,8 +739,14 @@ impl ListingDatabase {
 
     /// Inherit storage options from the connection into the target map
     fn inherit_storage_options(&self, target: &mut HashMap<String, String>) {
+        // Credential precedence applies to the whole family, not individual members. Once an
+        // operation supplies any member, the lower-precedence connection family must not fill in
+        // its missing token (or key/secret); the provider boundary validates completeness later.
+        let operation_has_aws_credentials = target.keys().any(|key| is_aws_credential_option(key));
         for (key, value) in self.storage_options.iter() {
-            if !target.contains_key(key) {
+            if !target.contains_key(key)
+                && !(operation_has_aws_credentials && is_aws_credential_option(key))
+            {
                 target.insert(key.clone(), value.clone());
             }
         }
@@ -1371,6 +1377,71 @@ mod tests {
         );
     }
 
+    #[cfg(feature = "aws")]
+    #[tokio::test]
+    async fn operation_credential_family_does_not_inherit_connection_token() {
+        let tempdir = tempdir().unwrap();
+        let request = ConnectRequest {
+            uri: tempdir.path().to_string_lossy().into_owned(),
+            #[cfg(feature = "remote")]
+            client_config: Default::default(),
+            options: HashMap::from([
+                (
+                    "aws_access_key_id".to_string(),
+                    "connection-key".to_string(),
+                ),
+                (
+                    "aws_secret_access_key".to_string(),
+                    "connection-secret".to_string(),
+                ),
+                (
+                    "aws_session_token".to_string(),
+                    "connection-token".to_string(),
+                ),
+            ]),
+            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 create_request = CreateTableRequest {
+            name: "operation_credentials".to_string(),
+            namespace_path: vec![],
+            data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
+            mode: CreateTableMode::Create,
+            write_options: WriteOptions {
+                lance_write_params: Some(lance::dataset::WriteParams {
+                    store_params: Some(operation_aws_store_params()),
+                    ..Default::default()
+                }),
+            },
+            location: None,
+            namespace_client: None,
+        };
+
+        let write_params = db.prepare_write_params(&create_request, None, None, None);
+        let options = write_params
+            .store_params
+            .unwrap()
+            .storage_options()
+            .unwrap()
+            .clone();
+
+        assert_eq!(options.get("aws_access_key_id").unwrap(), "operation-key");
+        assert_eq!(
+            options.get("aws_secret_access_key").unwrap(),
+            "operation-secret"
+        );
+        assert!(
+            !options.contains_key("aws_session_token"),
+            "an explicit operation family must not inherit a lower-precedence token"
+        );
+    }
+
     #[tokio::test]
     async fn test_listing_database_root_ops_do_not_create_manifest() {
         let tempdir = tempdir().unwrap();

rust/lancedb/src/io/object_store.rs

@@ -20,7 +20,7 @@ use lance::io::{ObjectStoreParams, WrappingObjectStore};
 #[cfg(feature = "aws")]
 use lance_io::object_store::{
     ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, StorageOptions,
-    providers::aws::build_aws_credential,
+    providers::aws::{AwsStoreProvider, build_aws_credential},
     throttle::{AimdThrottleConfig, AimdThrottledStore},
 };
 use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
@@ -98,6 +98,17 @@ fn is_aws_credential_key(key: &AmazonS3ConfigKey) -> bool {
     )
 }
 
+#[cfg(feature = "aws")]
+pub(crate) fn is_aws_credential_option(key: &str) -> bool {
+    AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase())
+        .is_ok_and(|key| is_aws_credential_key(&key))
+}
+
+#[cfg(not(feature = "aws"))]
+pub(crate) fn is_aws_credential_option(_key: &str) -> bool {
+    false
+}
+
 #[cfg(feature = "aws")]
 fn canonical_noncredential_options(
     storage_options: &HashMap<String, String>,
@@ -410,9 +421,27 @@ struct AtomicAwsStoreProvider {
 }
 
 #[cfg(feature = "aws")]
-#[async_trait]
-impl ObjectStoreProvider for AtomicAwsStoreProvider {
-    async fn new_store(
+impl AtomicAwsStoreProvider {
+    const CACHE_GENERATION: &'static str = "lancedb-atomic-aws-v1";
+
+    /// Lance does not expose provider downcasting. Its built-in provider is a unit struct with a
+    /// stable derived Debug representation, which is the only available way to distinguish the
+    /// one provider whose OpenDAL implementation needs adaptation from an unknown custom one.
+    fn is_builtin_aws_provider(&self) -> bool {
+        format!("{:?}", self.inner.as_ref()) == format!("{:?}", AwsStoreProvider)
+    }
+
+    fn generated_prefix(
+        &self,
+        url: &url::Url,
+        storage_options: Option<&HashMap<String, String>>,
+    ) -> lance_core::Result<String> {
+        self.inner
+            .calculate_object_store_prefix(url, storage_options)
+            .map(|prefix| format!("{prefix}${}", Self::CACHE_GENERATION))
+    }
+
+    async fn new_store_inner(
         &self,
         base_path: url::Url,
         params: &ObjectStoreParams,
@@ -423,6 +452,13 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider {
             .is_some_and(|value| value == "true");
 
         if use_opendal {
+            // A registered custom provider owns its store construction contract. In particular,
+            // its returned store may add encryption, authorization, or wrapping behavior that
+            // cannot be reconstructed from ObjectStoreParams. Only adapt Lance's known built-in
+            // provider, whose OpenDAL path ignores both supported credential-provider fields.
+            if !self.is_builtin_aws_provider() {
+                return self.inner.new_store(base_path, params).await;
+            }
             if storage_options
                 .get("aws_provider_scheme")
                 .is_some_and(|scheme| !scheme.is_empty())
@@ -532,6 +568,23 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider {
         atomic_params.aws_credentials = Some(credential_provider);
         self.inner.new_store(base_path, &atomic_params).await
     }
+}
+
+#[cfg(feature = "aws")]
+#[async_trait]
+impl ObjectStoreProvider for AtomicAwsStoreProvider {
+    async fn new_store(
+        &self,
+        base_path: url::Url,
+        params: &ObjectStoreParams,
+    ) -> lance_core::Result<LanceObjectStore> {
+        let store_prefix = self.generated_prefix(&base_path, params.storage_options())?;
+        let mut store = self.new_store_inner(base_path, params).await?;
+        // The registry cache and the returned store must use the same identity. Lance compares
+        // these values when resolving external blob bases.
+        store.store_prefix = store_prefix;
+        Ok(store)
+    }
 
     fn extract_path(&self, url: &url::Url) -> lance_core::Result<Path> {
         self.inner.extract_path(url)
@@ -542,9 +595,7 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider {
         url: &url::Url,
         storage_options: Option<&HashMap<String, String>>,
     ) -> lance_core::Result<String> {
-        self.inner
-            .calculate_object_store_prefix(url, storage_options)
-            .map(|prefix| format!("{prefix}$lancedb-atomic-aws-v1"))
+        self.generated_prefix(url, storage_options)
     }
 }
 
@@ -869,6 +920,24 @@ mod credential_tests {
         }
     }
 
+    #[derive(Debug)]
+    struct CustomStoreProvider {
+        marker: Arc<dyn ObjectStore>,
+    }
+
+    #[async_trait]
+    impl ObjectStoreProvider for CustomStoreProvider {
+        async fn new_store(
+            &self,
+            base_path: url::Url,
+            params: &ObjectStoreParams,
+        ) -> lance_core::Result<LanceObjectStore> {
+            let mut store = AwsStoreProvider.new_store(base_path, params).await?;
+            store.inner = self.marker.clone();
+            Ok(store)
+        }
+    }
+
     #[derive(Debug, PartialEq, Eq)]
     struct ObservedCredential {
         key_id: String,
@@ -1266,6 +1335,48 @@ mod credential_tests {
         );
     }
 
+    #[tokio::test]
+    async fn opendal_preserves_custom_provider_store_behavior() {
+        let marker: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
+        let provider = AtomicAwsStoreProvider {
+            inner: Arc::new(CustomStoreProvider {
+                marker: marker.clone(),
+            }),
+        };
+        let mut options = local_s3_options();
+        options.insert("use_opendal".to_string(), "true".to_string());
+
+        let store = provider
+            .new_store(
+                url::Url::parse("s3://bucket/table").unwrap(),
+                &object_store_params_from_storage_options(options),
+            )
+            .await
+            .unwrap();
+
+        assert!(
+            Arc::ptr_eq(&store.inner, &marker),
+            "the wrapper must not discard custom provider store behavior"
+        );
+    }
+
+    #[tokio::test]
+    async fn wrapper_store_prefix_matches_registry_identity() {
+        let registry = Arc::new(ObjectStoreRegistry::default());
+        let session = lance::session::Session::new(16, 16, registry.clone());
+        install_atomic_aws_provider(&session);
+        let uri = "s3://bucket/table";
+        let url = url::Url::parse(uri).unwrap();
+        let params = object_store_params_from_storage_options(local_s3_options());
+
+        let store = registry.get_store(url, &params).await.unwrap();
+        let registry_prefix = registry
+            .calculate_object_store_prefix(uri, params.storage_options())
+            .unwrap();
+
+        assert_eq!(store.store_prefix, registry_prefix);
+    }
+
     #[tokio::test]
     async fn identical_explicit_options_reuse_the_session_store() {
         let registry = Arc::new(ObjectStoreRegistry::default());

rust/lancedb/src/table.rs

@@ -2256,7 +2256,13 @@ impl NativeTable {
         pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
         managed_versioning: Option<bool>,
     ) -> Result<Self> {
-        let params = params.unwrap_or_default();
+        let mut params = params.unwrap_or_default();
+        let effective_session = params
+            .session
+            .clone()
+            .unwrap_or_else(|| Arc::new(lance::session::Session::default()));
+        crate::io::object_store::install_atomic_aws_provider(&effective_session);
+        params.session(effective_session);
         // patch the params if we have a write store wrapper
         let params = match write_store_wrapper.clone() {
             Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
@@ -2530,9 +2536,15 @@ impl NativeTable {
         pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
     ) -> Result<Self> {
         // Default params uses format v1.
-        let params = params.unwrap_or(WriteParams {
+        let mut params = params.unwrap_or(WriteParams {
             ..Default::default()
         });
+        let effective_session = params
+            .session
+            .clone()
+            .unwrap_or_else(|| Arc::new(lance::session::Session::default()));
+        crate::io::object_store::install_atomic_aws_provider(&effective_session);
+        params.session = Some(effective_session);
         // patch the params if we have a write store wrapper
         let params = match write_store_wrapper.clone() {
             Some(wrapper) => params.patch_with_store_wrapper(wrapper)?,
@@ -3581,6 +3593,10 @@ mod tests {
     use lance::Dataset;
     use lance::io::{ObjectStoreParams, WrappingObjectStore};
     use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
+    #[cfg(feature = "aws")]
+    use lance_io::object_store::{
+        ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry,
+    };
     use tempfile::tempdir;
 
     use super::*;
@@ -3590,6 +3606,53 @@ mod tests {
     use crate::query::{ExecutableQuery, QueryBase};
     use crate::test_utils::connection::new_test_connection;
 
+    #[cfg(feature = "aws")]
+    #[derive(Debug)]
+    struct RecordingS3Provider {
+        saw_atomic_credentials: Arc<AtomicBool>,
+    }
+
+    #[cfg(feature = "aws")]
+    #[async_trait]
+    impl ObjectStoreProvider for RecordingS3Provider {
+        async fn new_store(
+            &self,
+            _base_path: url::Url,
+            params: &ObjectStoreParams,
+        ) -> lance_core::Result<LanceObjectStore> {
+            self.saw_atomic_credentials
+                .store(params.aws_credentials.is_some(), Ordering::SeqCst);
+            Err(lance_core::Error::invalid_input("recorded test request"))
+        }
+    }
+
+    #[cfg(feature = "aws")]
+    fn recording_s3_session() -> (Arc<lance::session::Session>, Arc<AtomicBool>) {
+        let saw_atomic_credentials = Arc::new(AtomicBool::new(false));
+        let registry = Arc::new(ObjectStoreRegistry::default());
+        registry.insert(
+            "s3",
+            Arc::new(RecordingS3Provider {
+                saw_atomic_credentials: saw_atomic_credentials.clone(),
+            }),
+        );
+        (
+            Arc::new(lance::session::Session::new(16, 16, registry)),
+            saw_atomic_credentials,
+        )
+    }
+
+    #[cfg(feature = "aws")]
+    fn explicit_s3_store_params() -> ObjectStoreParams {
+        crate::io::object_store::object_store_params_from_storage_options(HashMap::from([
+            ("aws_access_key_id".to_string(), "explicit-key".to_string()),
+            (
+                "aws_secret_access_key".to_string(),
+                "explicit-secret".to_string(),
+            ),
+        ]))
+    }
+
     #[test]
     fn test_tokenize_uses_explicit_simple_tokenizer() {
         let params =
@@ -3651,6 +3714,70 @@ mod tests {
         );
     }
 
+    #[cfg(feature = "aws")]
+    #[tokio::test]
+    async fn direct_native_open_installs_the_atomic_provider() {
+        let (session, saw_atomic_credentials) = recording_s3_session();
+        let params = ReadParams {
+            session: Some(session),
+            store_options: Some(explicit_s3_store_params()),
+            ..Default::default()
+        };
+
+        let error = NativeTable::open_with_params(
+            "s3://bucket/table",
+            "table",
+            vec![],
+            None,
+            Some(params),
+            None,
+            None,
+            HashSet::new(),
+            None,
+        )
+        .await
+        .unwrap_err();
+
+        assert!(error.to_string().contains("recorded test request"));
+        assert!(
+            saw_atomic_credentials.load(Ordering::SeqCst),
+            "the public direct open path must install the wrapper on its session"
+        );
+    }
+
+    #[cfg(feature = "aws")]
+    #[tokio::test]
+    async fn direct_native_create_installs_the_atomic_provider() {
+        let (session, saw_atomic_credentials) = recording_s3_session();
+        let params = WriteParams {
+            session: Some(session),
+            store_params: Some(explicit_s3_store_params()),
+            ..Default::default()
+        };
+        let batch = make_test_batches();
+        let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
+
+        let error = NativeTable::create(
+            "s3://bucket/table",
+            "table",
+            vec![],
+            reader,
+            None,
+            Some(params),
+            None,
+            None,
+            HashSet::new(),
+        )
+        .await
+        .unwrap_err();
+
+        assert!(error.to_string().contains("recorded test request"));
+        assert!(
+            saw_atomic_credentials.load(Ordering::SeqCst),
+            "the public direct create path must install the wrapper on its session"
+        );
+    }
+
     /// Write a table and then break it, leaving the `<name>.lance` directory in place.
     ///
     /// `remove_all` reproduces an interrupted drop + re-create (the directory is left