builder, registry, split by schemes

Author: CTTY

crates/iceberg/src/io/file_io.rs

@@ -25,10 +25,9 @@ use async_trait::async_trait;
 use bytes::Bytes;
 use url::Url;
 
-use super::storage::OpenDALStorage;
 use crate::{Error, ErrorKind, Result};
 
-/// todo doc
+/// Trait for storage operations in Iceberg
 #[async_trait]
 pub trait Storage: Debug + Send + Sync {
     /// Check if a file exists at the given path
@@ -62,6 +61,16 @@ pub trait Storage: Debug + Send + Sync {
     fn new_output(&self, path: &str) -> Result<OutputFile>;
 }
 
+/// Common interface for all storage builders.
+pub trait StorageBuilder: Debug + Send + Sync {
+    /// Create a new storage instance with the given properties and extensions.
+    fn build(
+        &self,
+        props: HashMap<String, String>,
+        extensions: Extensions,
+    ) -> Result<Arc<dyn Storage>>;
+}
+
 /// FileIO implementation, used to manipulate files in underlying storage.
 ///
 /// # Note
@@ -288,10 +297,19 @@ impl FileIOBuilder {
 
     /// Builds [`FileIO`].
     pub fn build(self) -> Result<FileIO> {
-        let storage = OpenDALStorage::build(self.clone())?;
+        // Use the scheme to determine the storage type
+        let scheme = self.scheme_str.clone().unwrap_or_default();
+        
+        // Create registry and get builder
+        let registry = crate::io::StorageBuilderRegistry::new();
+        let builder = registry.get_builder(scheme.as_str())?;
+        
+        // Build storage with props and extensions
+        let storage = builder.build(self.props.clone(), self.extensions.clone())?;
+
         Ok(FileIO {
             builder: self,
-            inner: Arc::new(storage),
+            inner: storage,
         })
     }
 }

crates/iceberg/src/io/mod.rs

@@ -68,8 +68,10 @@
 
 mod file_io;
 mod storage;
+mod storage_builder;
 
 pub use file_io::*;
+pub use storage_builder::StorageBuilderRegistry;
 pub(crate) mod object_cache;
 
 #[cfg(feature = "storage-azdls")]

crates/iceberg/src/io/storage_azdls.rs

@@ -18,11 +18,18 @@
 use std::collections::HashMap;
 use std::fmt::Display;
 use std::str::FromStr;
+use std::sync::Arc;
 
-use opendal::Configurator;
+use async_trait::async_trait;
+use bytes::Bytes;
 use opendal::services::AzdlsConfig;
+use opendal::{Configurator, Operator};
 use url::Url;
 
+use crate::io::{
+    Extensions, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage,
+    StorageBuilder,
+};
 use crate::{Error, ErrorKind, Result, ensure_data_valid};
 
 /// A connection string.
@@ -125,7 +132,7 @@ pub(crate) fn azdls_create_operator<'a>(
 ///   paths are expected to contain the `dfs` storage service.
 /// - `wasb[s]` is used to refer to files in Blob Storage directly; paths are
 ///   expected to contain the `blob` storage service.
-#[derive(Debug, PartialEq)]
+#[derive(Debug, Clone, PartialEq)]
 pub(crate) enum AzureStorageScheme {
     Abfs,
     Abfss,
@@ -597,3 +604,116 @@ mod tests {
         }
     }
 }
+
+/// Azure Data Lake Storage implementation using OpenDAL
+#[derive(Debug, Clone)]
+pub struct OpenDALAzdlsStorage {
+    /// Because Azdls accepts multiple possible schemes, we store the full
+    /// passed scheme here to later validate schemes passed via paths.
+    configured_scheme: AzureStorageScheme,
+    config: Arc<AzdlsConfig>,
+}
+
+impl OpenDALAzdlsStorage {
+    /// Creates operator from path.
+    fn create_operator<'a>(&self, path: &'a str) -> Result<(Operator, &'a str)> {
+        let (op, relative_path) = azdls_create_operator(path, &self.config, &self.configured_scheme)?;
+        let op = op.layer(opendal::layers::RetryLayer::new());
+        Ok((op, relative_path))
+    }
+}
+
+#[async_trait]
+impl Storage for OpenDALAzdlsStorage {
+    async fn exists(&self, path: &str) -> Result<bool> {
+        let (op, relative_path) = self.create_operator(path)?;
+        Ok(op.exists(relative_path).await?)
+    }
+
+    async fn metadata(&self, path: &str) -> Result<FileMetadata> {
+        let (op, relative_path) = self.create_operator(path)?;
+        let meta = op.stat(relative_path).await?;
+        Ok(FileMetadata {
+            size: meta.content_length(),
+        })
+    }
+
+    async fn read(&self, path: &str) -> Result<Bytes> {
+        let (op, relative_path) = self.create_operator(path)?;
+        Ok(op.read(relative_path).await?.to_bytes())
+    }
+
+    async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
+        let (op, relative_path) = self.create_operator(path)?;
+        Ok(Box::new(op.reader(relative_path).await?))
+    }
+
+    async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
+        let mut writer = self.writer(path).await?;
+        writer.write(bs).await?;
+        writer.close().await
+    }
+
+    async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
+        let (op, relative_path) = self.create_operator(path)?;
+        Ok(Box::new(op.writer(relative_path).await?))
+    }
+
+    async fn delete(&self, path: &str) -> Result<()> {
+        let (op, relative_path) = self.create_operator(path)?;
+        Ok(op.delete(relative_path).await?)
+    }
+
+    async fn remove_dir_all(&self, path: &str) -> Result<()> {
+        let (op, relative_path) = self.create_operator(path)?;
+        let path = if relative_path.ends_with('/') {
+            relative_path.to_string()
+        } else {
+            format!("{relative_path}/")
+        };
+        Ok(op.remove_all(&path).await?)
+    }
+
+    fn new_input(&self, path: &str) -> Result<InputFile> {
+        Ok(InputFile {
+            storage: Arc::new(self.clone()),
+            path: path.to_string(),
+        })
+    }
+
+    fn new_output(&self, path: &str) -> Result<OutputFile> {
+        Ok(OutputFile {
+            storage: Arc::new(self.clone()),
+            path: path.to_string(),
+        })
+    }
+}
+
+/// Builder for Azure Data Lake Storage
+#[derive(Debug)]
+pub struct OpenDALAzdlsStorageBuilder;
+
+impl StorageBuilder for OpenDALAzdlsStorageBuilder {
+    fn build(
+        &self,
+        props: HashMap<String, String>,
+        _extensions: Extensions,
+    ) -> Result<Arc<dyn Storage>> {
+        // Get the scheme string from the props or use default
+        let scheme_str = props
+            .get("scheme_str")
+            .cloned()
+            .unwrap_or_else(|| "abfs".to_string());
+
+        // Parse the scheme
+        let scheme = scheme_str.parse::<AzureStorageScheme>()?;
+
+        // Parse Azdls config from props
+        let config = azdls_config_parse(props)?;
+
+        Ok(Arc::new(OpenDALAzdlsStorage {
+            configured_scheme: scheme,
+            config: Arc::new(config),
+        }))
+    }
+}