From 2a638a021bb6ded61dd275721dbb52f6e250a197 Mon Sep 17 00:00:00 2001 From: zhengpeng <847850277@qq.com> Date: Tue, 28 Jul 2026 14:58:41 +0800 Subject: [PATCH 1/2] fix: preserve ABFS container in object store registry keys --- datafusion/execution/src/object_store.rs | 59 +++++++++++++++++++++--- 1 file changed, 52 insertions(+), 7 deletions(-) diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index 22ce1f0cf2bbf..eb1cd5b1e5e9a 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -263,14 +263,18 @@ impl ObjectStoreRegistry for DefaultObjectStoreRegistry { } } -/// Get the key of a url for object store registration. -/// The credential info will be removed +/// Get the key of a URL for object store registration. +/// +/// Userinfo is preserved for ABFS schemes, where it identifies a namespace, +/// and removed for all other schemes. fn get_url_key(url: &Url) -> String { - format!( - "{}://{}", - url.scheme(), - &url[url::Position::BeforeHost..url::Position::AfterPort], - ) + let key_authority = match url.scheme() { + // ABFS encodes the container namespace in URL userinfo. + "abfs" | "abfss" => &url[url::Position::BeforeUsername..url::Position::AfterPort], + _ => &url[url::Position::BeforeHost..url::Position::AfterPort], + }; + + format!("{}://{key_authority}", url.scheme()) } #[cfg(test)] @@ -330,5 +334,46 @@ mod tests { let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap(); let key = get_url_key(&url.url); assert_eq!(key.as_str(), "s3://host:123"); + + for scheme in ["abfs", "abfss"] { + let url = ObjectStoreUrl::parse(format!( + "{scheme}://container@account.dfs.core.windows.net" + )) + .unwrap(); + let key = get_url_key(&url.url); + assert_eq!( + key, + format!("{scheme}://container@account.dfs.core.windows.net") + ); + } + } + + #[test] + fn test_abfs_containers_are_registered_separately() { + use object_store::memory::InMemory; + + for scheme in ["abfs", "abfss"] { + let registry = DefaultObjectStoreRegistry::new(); + let url_c1 = Url::parse(&format!( + "{scheme}://container1@account.dfs.core.windows.net/" + )) + .unwrap(); + let url_c2 = Url::parse(&format!( + "{scheme}://container2@account.dfs.core.windows.net/" + )) + .unwrap(); + + let store_c1: Arc = Arc::new(InMemory::new()); + let store_c2: Arc = Arc::new(InMemory::new()); + + registry.register_store(&url_c1, Arc::clone(&store_c1)); + registry.register_store(&url_c2, Arc::clone(&store_c2)); + + let actual_c1 = registry.get_store(&url_c1).unwrap(); + let actual_c2 = registry.get_store(&url_c2).unwrap(); + + assert!(Arc::ptr_eq(&actual_c1, &store_c1)); + assert!(Arc::ptr_eq(&actual_c2, &store_c2)); + } } } From fec42b43d0be74af1dc3608f20a872015530e238 Mon Sep 17 00:00:00 2001 From: zhengpeng <847850277@qq.com> Date: Tue, 28 Jul 2026 15:29:16 +0800 Subject: [PATCH 2/2] fix: preserve ABFS container in object store registry keys --- datafusion/execution/src/object_store.rs | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index eb1cd5b1e5e9a..f7d414f579af9 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -265,13 +265,17 @@ impl ObjectStoreRegistry for DefaultObjectStoreRegistry { /// Get the key of a URL for object store registration. /// -/// Userinfo is preserved for ABFS schemes, where it identifies a namespace, -/// and removed for all other schemes. +/// The username portion of userinfo is preserved for ABFS schemes, where it +/// identifies a namespace. Passwords and userinfo for other schemes are removed. fn get_url_key(url: &Url) -> String { let key_authority = match url.scheme() { // ABFS encodes the container namespace in URL userinfo. - "abfs" | "abfss" => &url[url::Position::BeforeUsername..url::Position::AfterPort], - _ => &url[url::Position::BeforeHost..url::Position::AfterPort], + "abfs" | "abfss" if !url.username().is_empty() => format!( + "{}@{}", + &url[url::Position::BeforeUsername..url::Position::AfterUsername], + &url[url::Position::BeforeHost..url::Position::AfterPort], + ), + _ => url[url::Position::BeforeHost..url::Position::AfterPort].to_string(), }; format!("{}://{key_authority}", url.scheme()) @@ -345,6 +349,16 @@ mod tests { key, format!("{scheme}://container@account.dfs.core.windows.net") ); + + let url = ObjectStoreUrl::parse(format!( + "{scheme}://container:secret@account.dfs.core.windows.net" + )) + .unwrap(); + let key = get_url_key(&url.url); + assert_eq!( + key, + format!("{scheme}://container@account.dfs.core.windows.net") + ); } }