Add support for Microsoft Azure in ObjectStore

Author: lizardoluis

Cargo.lock

@@ -37,7 +37,7 @@ getrandom = { version = "0.3.1", features = ["std"] }
 itertools = "0.14.0"
 lazy_static = "1.5.0"
 lru = "0.16.0"
-object_store = { version = "0.12", features = ["aws", "gcp"] }
+object_store = { version = "0.12", features = ["aws", "gcp", "azure"] }
 murmur3 = { version = "0.5.2" }
 parquet = { version = "56", features = ["async", "object_store"] }
 pin-project-lite = "0.2"

Cargo.toml

@@ -6,6 +6,7 @@ use std::{fmt::Display, path::Path, str::FromStr, sync::Arc};
 
 use object_store::{
     aws::{AmazonS3Builder, AmazonS3ConfigKey, S3CopyIfNotExists},
+    azure::{AzureConfigKey, MicrosoftAzureBuilder},
     gcp::{GoogleCloudStorageBuilder, GoogleConfigKey},
     local::LocalFileSystem,
     memory::InMemory,
@@ -24,6 +25,8 @@ pub enum Bucket<'s> {
     S3(&'s str),
     /// GCS bucket
     GCS(&'s str),
+    /// Azure container
+    Azure(&'s str),
     /// No bucket
     Local,
 }
@@ -33,13 +36,14 @@ impl Display for Bucket<'_> {
         match self {
             Bucket::S3(s) => write!(f, "s3://{s}"),
             Bucket::GCS(s) => write!(f, "gs://{s}"),
+            Bucket::Azure(s) => write!(f, "https://{s}"),
             Bucket::Local => write!(f, ""),
         }
     }
 }
 
 impl Bucket<'_> {
-    /// Get the bucket and coud provider from the location string
+    /// Get the bucket and cloud provider from the location string
     pub fn from_path(path: &str) -> Result<Bucket<'_>, Error> {
         if path.starts_with("s3://") || path.starts_with("s3a://") {
             let prefix = if path.starts_with("s3://") {
@@ -63,6 +67,17 @@ impl Bucket<'_> {
                 .next()
                 .map(Bucket::GCS)
                 .ok_or(Error::NotFound(format!("Bucket in path {path}")))
+        } else if path.starts_with("https://")
+            && (path.contains("dfs.core.windows.net")
+                || path.contains("blob.core.windows.net")
+                || path.contains("dfs.fabric.microsoft.com")
+                || path.contains("blob.fabric.microsoft.com"))
+        {
+            path.trim_start_matches("https://")
+                .split('/')
+                .nth(1)
+                .map(Bucket::Azure)
+                .ok_or(Error::NotFound(format!("Bucket in path {path}")))
         } else {
             Ok(Bucket::Local)
         }
@@ -72,6 +87,8 @@ impl Bucket<'_> {
 /// A wrapper for ObjectStore builders that can be used as a template to generate an ObjectStore given a particular bucket.
 #[derive(Debug, Clone)]
 pub enum ObjectStoreBuilder {
+    /// Microsoft Azure builder
+    Azure(Box<MicrosoftAzureBuilder>),
     /// AWS s3 builder
     S3(Box<AmazonS3Builder>),
     /// Google Cloud Storage builder
@@ -84,6 +101,8 @@ pub enum ObjectStoreBuilder {
 
 /// Configuration keys for [ObjectStoreBuilder]
 pub enum ConfigKey {
+    /// Configuration keys for Microsoft Azure
+    Azure(AzureConfigKey),
     /// Configuration keys for AWS S3
     AWS(AmazonS3ConfigKey),
     /// Configuration keys for GCS
@@ -93,6 +112,9 @@ pub enum ConfigKey {
 impl FromStr for ConfigKey {
     type Err = object_store::Error;
     fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Ok(x) = s.parse() {
+            return Ok(ConfigKey::Azure(x));
+        };
         if let Ok(x) = s.parse() {
             return Ok(ConfigKey::AWS(x));
         };
@@ -106,6 +128,10 @@ impl FromStr for ConfigKey {
     }
 }
 impl ObjectStoreBuilder {
+    /// Create a new Microsoft Azure ObjectStoreBuilder
+    pub fn azure() -> Self {
+        ObjectStoreBuilder::Azure(Box::new(MicrosoftAzureBuilder::from_env()))
+    }
     /// Create new AWS S3 Object Store builder
     pub fn s3() -> Self {
         ObjectStoreBuilder::S3(Box::new(AmazonS3Builder::from_env()))
@@ -125,6 +151,9 @@ impl ObjectStoreBuilder {
     /// Set config value for builder
     pub fn with_config(self, key: ConfigKey, value: impl Into<String>) -> Self {
         match (self, key) {
+            (ObjectStoreBuilder::Azure(azure), ConfigKey::Azure(key)) => {
+                ObjectStoreBuilder::Azure(Box::new(azure.with_config(key, value)))
+            }
             (ObjectStoreBuilder::S3(aws), ConfigKey::AWS(key)) => {
                 ObjectStoreBuilder::S3(Box::new(aws.with_config(key, value)))
             }
@@ -137,6 +166,13 @@ impl ObjectStoreBuilder {
     /// Create objectstore from template
     pub fn build(&self, bucket: Bucket) -> Result<Arc<dyn ObjectStore>, Error> {
         match (bucket, self) {
+            (Bucket::Azure(bucket), Self::Azure(builder)) => Ok::<_, Error>(Arc::new(
+                (**builder)
+                    .clone()
+                    .with_container_name(bucket)
+                    .build()
+                    .map_err(Error::from)?,
+            )),
             (Bucket::S3(bucket), Self::S3(builder)) => Ok::<_, Error>(Arc::new(
                 (**builder)
                     .clone()
@@ -158,3 +194,106 @@ impl ObjectStoreBuilder {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_from_path_s3() {
+        let bucket = Bucket::from_path("s3://my-bucket/path/to/file").unwrap();
+        match bucket {
+            Bucket::S3(name) => assert_eq!(name, "my-bucket"),
+            _ => panic!("Expected S3 bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_s3a() {
+        let bucket = Bucket::from_path("s3a://my-bucket/path/to/file").unwrap();
+        match bucket {
+            Bucket::S3(name) => assert_eq!(name, "my-bucket"),
+            _ => panic!("Expected S3 bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_gcs() {
+        let bucket = Bucket::from_path("gcs://my-bucket/path/to/file").unwrap();
+        match bucket {
+            Bucket::GCS(name) => assert_eq!(name, "my-bucket"),
+            _ => panic!("Expected GCS bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_gs() {
+        let bucket = Bucket::from_path("gs://my-bucket/path/to/file").unwrap();
+        match bucket {
+            Bucket::GCS(name) => assert_eq!(name, "my-bucket"),
+            _ => panic!("Expected GCS bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_azure_dfs() {
+        let bucket =
+            Bucket::from_path("https://mystorageaccount.dfs.core.windows.net/container/path")
+                .unwrap();
+        match bucket {
+            Bucket::Azure(name) => assert_eq!(name, "container"),
+            _ => panic!("Expected Azure bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_azure_blob() {
+        let bucket =
+            Bucket::from_path("https://mystorageaccount.blob.core.windows.net/container/path")
+                .unwrap();
+        match bucket {
+            Bucket::Azure(name) => assert_eq!(name, "container"),
+            _ => panic!("Expected Azure bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_azure_fabric_dfs() {
+        let bucket =
+            Bucket::from_path("https://mystorageaccount.dfs.fabric.microsoft.com/container/path")
+                .unwrap();
+        match bucket {
+            Bucket::Azure(name) => assert_eq!(name, "container"),
+            _ => panic!("Expected Azure bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_azure_fabric_blob() {
+        let bucket =
+            Bucket::from_path("https://mystorageaccount.blob.fabric.microsoft.com/container/path")
+                .unwrap();
+        match bucket {
+            Bucket::Azure(name) => assert_eq!(name, "container"),
+            _ => panic!("Expected Azure bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_local() {
+        let bucket = Bucket::from_path("/local/path/to/file").unwrap();
+        match bucket {
+            Bucket::Local => {}
+            _ => panic!("Expected Local bucket"),
+        }
+    }
+
+    #[test]
+    fn test_from_path_https_non_azure() {
+        let bucket = Bucket::from_path("https://example.com/path").unwrap();
+        match bucket {
+            Bucket::Local => {}
+            _ => panic!("Expected Local bucket"),
+        }
+    }
+}

iceberg-rust/src/object_store/mod.rs

@@ -3,6 +3,7 @@
 
 use crate::error::Error;
 use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey};
+use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder};
 use object_store::gcp::{GcpCredential, GoogleCloudStorageBuilder, GoogleConfigKey};
 use object_store::{parse_url_opts, ObjectStore, ObjectStoreScheme, StaticCredentialProvider};
 use std::collections::HashMap;
@@ -23,6 +24,12 @@ const GCS_BUCKET: &str = "gcs.bucket";
 const GCS_CREDENTIALS_JSON: &str = "gcs.credentials-json";
 const GCS_TOKEN: &str = "gcs.oauth2.token";
 
+/// Azure configs
+const AZURE_CONTAINER_NAME: &str = "azure.container-name";
+const AZURE_ENDPOINT: &str = "azure.endpoint";
+const AZURE_STORAGE_ACCESS_KEY: &str = "azure.access-key";
+const AZURE_STORAGE_ACCOUNT_NAME: &str = "azure.account-name";
+
 /// Parse the url and Iceberg format of variuos storage options into the equivalent `object_store`
 /// options and build the corresponding `ObjectStore`.
 pub fn object_store_from_config(
@@ -73,6 +80,27 @@ pub fn object_store_from_config(
             Arc::new(builder.build()?) as Arc<dyn ObjectStore>
         }
 
+        (ObjectStoreScheme::MicrosoftAzure, _) => {
+            let mut builder = MicrosoftAzureBuilder::new().with_url(url);
+            for (key, option) in config {
+                let azure_key = match key.as_str() {
+                    AZURE_CONTAINER_NAME => AzureConfigKey::ContainerName,
+                    AZURE_STORAGE_ACCOUNT_NAME => AzureConfigKey::AccountName,
+                    AZURE_STORAGE_ACCESS_KEY => AzureConfigKey::AccessKey,
+                    AZURE_ENDPOINT => {
+                        if option.starts_with("http://") {
+                            // This is mainly used for testing, e.g. against Azurite
+                            builder = builder.with_allow_http(true);
+                        }
+                        AzureConfigKey::Endpoint
+                    }
+                    _ => continue,
+                };
+                builder = builder.with_config(azure_key, option);
+            }
+            Arc::new(builder.build()?) as Arc<dyn ObjectStore>
+        }
+
         _ => {
             let (store, _path) = parse_url_opts(&url, config)?;
             store.into()
@@ -169,4 +197,51 @@ mod tests {
         assert!(store_repr.contains("bearer: \"oauth-token-123\""));
         assert!(store_repr.contains("bucket_name: \"test-bucket\""));
     }
+
+    #[test]
+    fn test_azure_config_basic() {
+        let url = Url::parse("https://testaccount.blob.core.windows.net/test-container").unwrap();
+        let mut config = HashMap::new();
+        config.insert(
+            AZURE_STORAGE_ACCOUNT_NAME.to_string(),
+            "testaccount".to_string(),
+        );
+        config.insert(AZURE_STORAGE_ACCESS_KEY.to_string(), "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==".to_string());
+
+        let store = object_store_from_config(url, config).unwrap();
+        let store_repr = format!("{store:?}");
+
+        println!("{}", store_repr);
+        assert!(store_repr.contains("account: \"testaccount\""));
+        assert!(store_repr.contains("container: \"test-container\""));
+        assert!(store_repr.contains("host: Some(Domain(\"testaccount.blob.core.windows.net\"))"));
+        assert!(store_repr.contains("port: None"));
+        assert!(store_repr.contains("scheme: \"https\""));
+        assert!(store_repr.contains("allow_http: Parsed(false)"));
+    }
+
+    #[test]
+    fn test_azure_config_with_http_endpoint() {
+        let url = Url::parse("https://testaccount.blob.core.windows.net/test-container").unwrap();
+        let mut config = HashMap::new();
+        config.insert(
+            AZURE_ENDPOINT.to_string(),
+            "http://localhost:9000".to_string(),
+        );
+        config.insert(
+            AZURE_STORAGE_ACCOUNT_NAME.to_string(),
+            "testaccount".to_string(),
+        );
+        config.insert(AZURE_STORAGE_ACCESS_KEY.to_string(), "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==".to_string());
+
+        let store = object_store_from_config(url, config).unwrap();
+        let store_repr = format!("{store:?}");
+
+        assert!(store_repr.contains("account: \"testaccount\""));
+        assert!(store_repr.contains("container: \"test-container\""));
+        assert!(store_repr.contains("host: Some(Domain(\"localhost\"))"));
+        assert!(store_repr.contains("port: Some(9000)"));
+        assert!(store_repr.contains("scheme: \"http\""));
+        assert!(store_repr.contains("allow_http: Parsed(true)"));
+    }
 }