Add support for custom credential loader for S3 FileIO (#1528)

## Which issue does this PR close?

- Closes #1527

## What changes are included in this PR?

Adds the ability to provide custom extensions to the `FileIOBuilder`.
Currently the only supported extension is `CustomAwsCredentialLoader`
which is a newtype around
[`AwsCredentialLoad`](https://docs.rs/reqsign/0.16.3/reqsign/trait.AwsCredentialLoad.html),
which is what OpenDAL expects.

I've added extensions to the `RestCatalog` as well, and when its
constructing `FileIO` for table operations, passes along any defined
extensions into the `FileIOBuilder`, which then get passed into the
underlying OpenDAL constructors.

## Are these changes tested?

Yes, tests added in `crates/iceberg/tests/file_io_s3_test.rs` that
verify the extension is working.

Author: phillipleblanc

Cargo.lock

@@ -17,11 +17,12 @@
 
 //! This module contains the iceberg REST catalog implementation.
 
+use std::any::Any;
 use std::collections::HashMap;
 use std::str::FromStr;
 
 use async_trait::async_trait;
-use iceberg::io::FileIO;
+use iceberg::io::{self, FileIO};
 use iceberg::table::Table;
 use iceberg::{
     Catalog, Error, ErrorKind, Namespace, NamespaceIdent, Result, TableCommit, TableCreation,
@@ -240,6 +241,8 @@ pub struct RestCatalog {
     /// It's could be different from the config fetched from the server and used at runtime.
     user_config: RestCatalogConfig,
     ctx: OnceCell<RestContext>,
+    /// Extensions for the FileIOBuilder.
+    file_io_extensions: io::Extensions,
 }
 
 impl RestCatalog {
@@ -248,9 +251,16 @@ impl RestCatalog {
         Self {
             user_config: config,
             ctx: OnceCell::new(),
+            file_io_extensions: io::Extensions::default(),
         }
     }
 
+    /// Add an extension to the file IO builder.
+    pub fn with_file_io_extension<T: Any + Send + Sync>(mut self, ext: T) -> Self {
+        self.file_io_extensions.add(ext);
+        self
+    }
+
     /// Gets the [`RestContext`] from the catalog.
     async fn context(&self) -> Result<&RestContext> {
         self.ctx
@@ -307,7 +317,10 @@ impl RestCatalog {
         };
 
         let file_io = match warehouse_path.or(metadata_location) {
-            Some(url) => FileIO::from_path(url)?.with_props(props).build()?,
+            Some(url) => FileIO::from_path(url)?
+                .with_props(props)
+                .with_extensions(self.file_io_extensions.clone())
+                .build()?,
             None => {
                 return Err(Error::new(
                     ErrorKind::Unexpected,

crates/catalog/rest/src/catalog.rs

@@ -37,7 +37,7 @@ storage-fs = ["opendal/services-fs"]
 storage-gcs = ["opendal/services-gcs"]
 storage-memory = ["opendal/services-memory"]
 storage-oss = ["opendal/services-oss"]
-storage-s3 = ["opendal/services-s3"]
+storage-s3 = ["opendal/services-s3", "reqsign"]
 
 async-std = ["dep:async-std"]
 tokio = ["tokio/rt-multi-thread"]
@@ -76,6 +76,7 @@ ordered-float = { workspace = true }
 parquet = { workspace = true, features = ["async"] }
 rand = { workspace = true }
 reqwest = { workspace = true }
+reqsign = { version = "0.16.3", optional = true, default-features = false }
 roaring = { workspace = true }
 rust_decimal = { workspace = true }
 serde = { workspace = true }

crates/iceberg/Cargo.toml

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::any::{Any, TypeId};
 use std::collections::HashMap;
 use std::ops::Range;
 use std::sync::Arc;
@@ -167,6 +168,31 @@ impl FileIO {
     }
 }
 
+/// Container for storing type-safe extensions used to configure underlying FileIO behavior.
+#[derive(Clone, Debug, Default)]
+pub struct Extensions(HashMap<TypeId, Arc<dyn Any + Send + Sync>>);
+
+impl Extensions {
+    /// Add an extension.
+    pub fn add<T: Any + Send + Sync>(&mut self, ext: T) {
+        self.0.insert(TypeId::of::<T>(), Arc::new(ext));
+    }
+
+    /// Extends the current set of extensions with another set of extensions.
+    pub fn extend(&mut self, extensions: Extensions) {
+        self.0.extend(extensions.0);
+    }
+
+    /// Fetch an extension.
+    pub fn get<T>(&self) -> Option<Arc<T>>
+    where T: 'static + Send + Sync + Clone {
+        let type_id = TypeId::of::<T>();
+        self.0
+            .get(&type_id)
+            .and_then(|arc_any| Arc::clone(arc_any).downcast::<T>().ok())
+    }
+}
+
 /// Builder for [`FileIO`].
 #[derive(Clone, Debug)]
 pub struct FileIOBuilder {
@@ -176,6 +202,8 @@ pub struct FileIOBuilder {
     scheme_str: Option<String>,
     /// Arguments for operator.
     props: HashMap<String, String>,
+    /// Optional extensions to configure the underlying FileIO behavior.
+    extensions: Extensions,
 }
 
 impl FileIOBuilder {
@@ -185,6 +213,7 @@ impl FileIOBuilder {
         Self {
             scheme_str: Some(scheme_str.to_string()),
             props: HashMap::default(),
+            extensions: Extensions::default(),
         }
     }
 
@@ -193,14 +222,19 @@ impl FileIOBuilder {
         Self {
             scheme_str: None,
             props: HashMap::default(),
+            extensions: Extensions::default(),
         }
     }
 
     /// Fetch the scheme string.
     ///
     /// The scheme_str will be empty if it's None.
-    pub fn into_parts(self) -> (String, HashMap<String, String>) {
-        (self.scheme_str.unwrap_or_default(), self.props)
+    pub fn into_parts(self) -> (String, HashMap<String, String>, Extensions) {
+        (
+            self.scheme_str.unwrap_or_default(),
+            self.props,
+            self.extensions,
+        )
     }
 
     /// Add argument for operator.
@@ -219,6 +253,24 @@ impl FileIOBuilder {
         self
     }
 
+    /// Add an extension to the file IO builder.
+    pub fn with_extension<T: Any + Send + Sync>(mut self, ext: T) -> Self {
+        self.extensions.add(ext);
+        self
+    }
+
+    /// Adds multiple extensions to the file IO builder.
+    pub fn with_extensions(mut self, extensions: Extensions) -> Self {
+        self.extensions.extend(extensions);
+        self
+    }
+
+    /// Fetch an extension from the file IO builder.
+    pub fn extension<T>(&self) -> Option<Arc<T>>
+    where T: 'static + Send + Sync + Clone {
+        self.extensions.get::<T>()
+    }
+
     /// Builds [`FileIO`].
     pub fn build(self) -> Result<FileIO> {
         let storage = Storage::build(self.clone())?;

crates/iceberg/src/io/file_io.rs

@@ -31,6 +31,8 @@ use opendal::{Operator, Scheme};
 #[cfg(feature = "storage-azdls")]
 use super::AzureStorageScheme;
 use super::FileIOBuilder;
+#[cfg(feature = "storage-s3")]
+use crate::io::CustomAwsCredentialLoader;
 use crate::{Error, ErrorKind};
 
 /// The storage carries all supported storage services in iceberg
@@ -47,6 +49,7 @@ pub(crate) enum Storage {
         /// Storing the scheme string here to return the correct path.
         configured_scheme: String,
         config: Arc<S3Config>,
+        customized_credential_load: Option<CustomAwsCredentialLoader>,
     },
     #[cfg(feature = "storage-gcs")]
     Gcs { config: Arc<GcsConfig> },
@@ -67,7 +70,7 @@ pub(crate) enum Storage {
 impl Storage {
     /// Convert iceberg config to opendal config.
     pub(crate) fn build(file_io_builder: FileIOBuilder) -> crate::Result<Self> {
-        let (scheme_str, props) = file_io_builder.into_parts();
+        let (scheme_str, props, extensions) = file_io_builder.into_parts();
         let scheme = Self::parse_scheme(&scheme_str)?;
 
         match scheme {
@@ -79,6 +82,9 @@ impl Storage {
             Scheme::S3 => Ok(Self::S3 {
                 configured_scheme: scheme_str,
                 config: super::s3_config_parse(props)?.into(),
+                customized_credential_load: extensions
+                    .get::<CustomAwsCredentialLoader>()
+                    .map(Arc::unwrap_or_clone),
             }),
             #[cfg(feature = "storage-gcs")]
             Scheme::Gcs => Ok(Self::Gcs {
@@ -144,8 +150,9 @@ impl Storage {
             Storage::S3 {
                 configured_scheme,
                 config,
+                customized_credential_load,
             } => {
-                let op = super::s3_config_build(config, path)?;
+                let op = super::s3_config_build(config, customized_credential_load, path)?;
                 let op_info = op.info();
 
                 // Check prefix of s3 path.

crates/iceberg/src/io/storage.rs

@@ -16,9 +16,13 @@
 // under the License.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
+use async_trait::async_trait;
 use opendal::services::S3Config;
 use opendal::{Configurator, Operator};
+pub use reqsign::{AwsCredential, AwsCredentialLoad};
+use reqwest::Client;
 use url::Url;
 
 use crate::io::is_truthy;
@@ -151,7 +155,11 @@ pub(crate) fn s3_config_parse(mut m: HashMap<String, String>) -> Result<S3Config
 }
 
 /// Build new opendal operator from give path.
-pub(crate) fn s3_config_build(cfg: &S3Config, path: &str) -> Result<Operator> {
+pub(crate) fn s3_config_build(
+    cfg: &S3Config,
+    customized_credential_load: &Option<CustomAwsCredentialLoader>,
+    path: &str,
+) -> Result<Operator> {
     let url = Url::parse(path)?;
     let bucket = url.host_str().ok_or_else(|| {
         Error::new(
@@ -160,11 +168,49 @@ pub(crate) fn s3_config_build(cfg: &S3Config, path: &str) -> Result<Operator> {
         )
     })?;
 
-    let builder = cfg
+    let mut builder = cfg
         .clone()
         .into_builder()
         // Set bucket name.
         .bucket(bucket);
 
+    if let Some(customized_credential_load) = customized_credential_load {
+        builder = builder
+            .customized_credential_load(customized_credential_load.clone().into_opendal_loader());
+    }
+
     Ok(Operator::new(builder)?.finish())
 }
+
+/// Custom AWS credential loader.
+/// This can be used to load credentials from a custom source, such as the AWS SDK.
+///
+/// This should be set as an extension on `FileIOBuilder`.
+#[derive(Clone)]
+pub struct CustomAwsCredentialLoader(Arc<dyn AwsCredentialLoad>);
+
+impl std::fmt::Debug for CustomAwsCredentialLoader {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CustomAwsCredentialLoader")
+            .finish_non_exhaustive()
+    }
+}
+
+impl CustomAwsCredentialLoader {
+    /// Create a new custom AWS credential loader.
+    pub fn new(loader: Arc<dyn AwsCredentialLoad>) -> Self {
+        Self(loader)
+    }
+
+    /// Convert this loader into an opendal compatible loader for customized AWS credentials.
+    pub fn into_opendal_loader(self) -> Box<dyn AwsCredentialLoad> {
+        Box::new(self)
+    }
+}
+
+#[async_trait]
+impl AwsCredentialLoad for CustomAwsCredentialLoader {
+    async fn load_credential(&self, client: Client) -> anyhow::Result<Option<AwsCredential>> {
+        self.0.load_credential(client).await
+    }
+}