Pass scheme strings cleanly

Author: CTTY

crates/iceberg/src/io/file_io.rs

@@ -27,7 +27,7 @@ use url::Url;
 
 // Re-export traits from storage module
 pub use super::storage::{Storage, StorageBuilder};
-use crate::io::StorageBuilderRegistry;
+use crate::io::{StorageBuilderRegistry, STORAGE_LOCATION_SCHEME};
 use crate::{Error, ErrorKind, Result};
 
 /// FileIO implementation, used to manipulate files in underlying storage.
@@ -321,8 +321,11 @@ impl FileIOBuilder {
 
         let builder = registry.get_builder(scheme.as_str())?;
 
+        let mut props_with_scheme = self.props.clone();
+        props_with_scheme.insert(STORAGE_LOCATION_SCHEME.to_string(), scheme);
+
         // Build storage with props and extensions
-        let storage = builder.build(self.props.clone(), self.extensions.clone())?;
+        let storage = builder.build(props_with_scheme, self.extensions.clone())?;
 
         Ok(FileIO {
             builder: self,

crates/iceberg/src/io/mod.rs

@@ -73,6 +73,12 @@ pub use file_io::*;
 pub use storage::{Storage, StorageBuilder, StorageBuilderRegistry};
 pub(crate) mod object_cache;
 
+/// Property key used to pass the scheme string to storage builders.
+/// 
+/// Custom storage implementations can use this constant to retrieve the scheme
+/// from the properties map passed to `StorageBuilder::build()`.
+pub const STORAGE_LOCATION_SCHEME: &str = "iceberg.storage.location.scheme";
+
 #[cfg(feature = "storage-azdls")]
 mod storage_azdls;
 #[cfg(feature = "storage-fs")]

crates/iceberg/src/io/storage_azdls.rs

@@ -28,6 +28,7 @@ use url::Url;
 
 use crate::io::{
     Extensions, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageBuilder,
+    STORAGE_LOCATION_SCHEME,
 };
 use crate::{Error, ErrorKind, Result, ensure_data_valid};
 
@@ -695,7 +696,7 @@ impl StorageBuilder for OpenDALAzdlsStorageBuilder {
     ) -> Result<Arc<dyn Storage>> {
         // Get the scheme string from the props or use default
         let scheme_str = props
-            .get("scheme_str")
+            .get(STORAGE_LOCATION_SCHEME)
             .cloned()
             .unwrap_or_else(|| "abfs".to_string());
 

crates/iceberg/src/io/storage_s3.rs

@@ -28,7 +28,7 @@ use url::Url;
 
 use crate::io::{
     Extensions, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageBuilder,
-    is_truthy,
+    is_truthy, STORAGE_LOCATION_SCHEME,
 };
 use crate::{Error, ErrorKind, Result};
 
@@ -334,7 +334,7 @@ impl StorageBuilder for OpenDALS3StorageBuilder {
     ) -> Result<Arc<dyn Storage>> {
         // Get the scheme string from the props or use "s3" as default
         let scheme_str = props
-            .get("scheme_str")
+            .get(STORAGE_LOCATION_SCHEME)
             .cloned()
             .unwrap_or_else(|| "s3".to_string());