feat: Infer storage name based on endpoint

Cargo.toml

@@ -105,6 +105,7 @@ backon = "0.4.0"
 base64 = "0.21"
 bb8 = { version = "0.8", optional = true }
 bytes = "1.2"
+ctor = "0.1.3"
 dashmap = { version = "5.4", optional = true }
 flagset = "0.4"
 futures = { version = "0.3", features = ["alloc"] }
@@ -131,6 +132,7 @@ redis = { version = "0.22", features = [
   "tokio-comp",
   "connection-manager",
 ], optional = true }
+regex = "1.7.1"
 reqsign = "0.8.3"
 reqwest = { version = "0.11.13", features = [
   "multipart",

src/services/azblob/backend.rs

@@ -16,17 +16,18 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 use std::fmt::Formatter;
 use std::fmt::Write;
-use std::mem;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use ctor::ctor;
 use http::header::HeaderName;
 use http::header::CONTENT_LENGTH;
 use http::header::CONTENT_TYPE;
 use http::Request;
 use http::Response;
 use http::StatusCode;
 use log::debug;
+use regex::Regex;
 use reqsign::AzureStorageSigner;
 
 use super::error::parse_error;
@@ -39,6 +40,15 @@ use crate::*;
 
 const X_MS_BLOB_TYPE: &str = "x-ms-blob-type";
 
+/// Known regex pattern Azure Storage Blob services resource URI syntax.
+/// Azure public cloud: https://accountname.blob.core.windows.net
+/// Azure US Government: https://accountname.blob.core.usgovcloudapi.net
+/// Azure China: https://accountname.blob.core.chinacloudapi.cn
+#[ctor]
+static KNOWN_AZBLOB_RESOURCE_URI_SYNTAX_REGEX: Regex =
+    Regex::new(r"(?i)^https?://([a-z0-9]{3,24})\.blob\.core\.(?:usgovcloudapi\.net|chinacloudapi\.cn|windows\.net)/?$")
+    .unwrap();
+
 /// Azure Storage Blob services support.
 ///
 /// # Capabilities
@@ -366,10 +376,23 @@ impl Builder for AzblobBuilder {
         };
 
         let mut signer_builder = AzureStorageSigner::builder();
+        let mut account_name: Option<String> = None;
         if let Some(sas_token) = &self.sas_token {
             signer_builder.security_token(sas_token);
+            match &self.account_name {
+                Some(name) => account_name = Some(name.clone()),
+                None => {
+                    account_name = infer_storage_name_from_endpoint(endpoint.as_str());
+                }
+            }
         } else if let (Some(name), Some(key)) = (&self.account_name, &self.account_key) {
+            account_name = Some(name.clone());
             signer_builder.account_name(name).account_key(key);
+        } else if let Some(key) = &self.account_key {
+            account_name = infer_storage_name_from_endpoint(endpoint.as_str());
+            signer_builder
+                .account_name(account_name.as_ref().unwrap_or(&String::new()))
+                .account_key(key);
         }
 
         let signer = signer_builder.build().map_err(|e| {
@@ -388,11 +411,27 @@ impl Builder for AzblobBuilder {
             signer: Arc::new(signer),
             container: self.container.clone(),
             client,
-            _account_name: mem::take(&mut self.account_name).unwrap_or_default(),
+            _account_name: account_name.unwrap_or_else(String::new),
         })
     }
 }
 
+fn infer_storage_name_from_endpoint(endpoint: &str) -> Option<String> {
+    let candidate_names = KNOWN_AZBLOB_RESOURCE_URI_SYNTAX_REGEX.captures(endpoint);
+    if candidate_names.is_none() || candidate_names.as_ref().unwrap().len() != 2 {
+        return None;
+    }
+    Some(
+        candidate_names
+            .unwrap()
+            .get(1)
+            .unwrap()
+            .clone()
+            .as_str()
+            .to_string(),
+    )
+}
+
 /// Backend for azblob services.
 #[derive(Debug, Clone)]
 pub struct AzblobBackend {
@@ -685,8 +724,99 @@ impl AzblobBackend {
 
 #[cfg(test)]
 mod tests {
+    use crate::{services::azblob::backend::infer_storage_name_from_endpoint, Builder};
+
     use super::AzblobBuilder;
 
+    #[test]
+    fn test_infer_storage_name_from_endpoint() {
+        let endpoint = "https://account.blob.core.windows.net";
+        let storage_name = infer_storage_name_from_endpoint(endpoint);
+        assert_eq!(storage_name, Some("account".to_string()));
+    }
+
+    #[test]
+    fn test_infer_storage_name_from_endpoint_with_trailing_slash() {
+        let endpoint = "https://account.blob.core.windows.net/";
+        let storage_name = infer_storage_name_from_endpoint(endpoint);
+        assert_eq!(storage_name, Some("account".to_string()));
+    }
+
+    #[test]
+    fn test_builder_from_endpoint_and_key_infer_account_name() {
+        let mut azblob_builder = AzblobBuilder::default();
+        azblob_builder.endpoint("https://storagesample.blob.core.chinacloudapi.cn");
+        azblob_builder.container("container");
+        azblob_builder.account_key("account-key");
+        let azblob = azblob_builder
+            .build()
+            .expect("build azblob should be succeeded.");
+
+        assert_eq!(
+            azblob.endpoint,
+            "https://storagesample.blob.core.chinacloudapi.cn"
+        );
+
+        assert_eq!(azblob._account_name, "storagesample".to_string());
+
+        assert_eq!(azblob.container, "container".to_string());
+
+        assert_eq!(
+            azblob_builder.account_key.unwrap(),
+            "account-key".to_string()
+        );
+    }
+
+    #[test]
+    fn test_no_key_wont_infer_account_name() {
+        let mut azblob_builder = AzblobBuilder::default();
+        azblob_builder.endpoint("https://storagesample.blob.core.windows.net");
+        azblob_builder.container("container");
+        let azblob = azblob_builder
+            .build()
+            .expect("build azblob should be succeeded.");
+
+        assert_eq!(
+            azblob.endpoint,
+            "https://storagesample.blob.core.windows.net"
+        );
+
+        assert_eq!(azblob._account_name, "".to_string());
+
+        assert_eq!(azblob.container, "container".to_string());
+
+        assert_eq!(azblob_builder.account_key, None);
+    }
+
+    #[test]
+    fn test_builder_from_endpoint_and_sas() {
+        let mut azblob_builder = AzblobBuilder::default();
+        azblob_builder.endpoint("https://storagesample.blob.core.usgovcloudapi.net");
+        azblob_builder.container("container");
+        azblob_builder.account_name("storagesample");
+        azblob_builder.account_key("account-key");
+        azblob_builder.sas_token("sas");
+        let azblob = azblob_builder
+            .build()
+            .expect("build azblob should be succeeded.");
+
+        assert_eq!(
+            azblob.endpoint,
+            "https://storagesample.blob.core.usgovcloudapi.net"
+        );
+
+        assert_eq!(azblob._account_name, "storagesample".to_string());
+
+        assert_eq!(azblob.container, "container".to_string());
+
+        assert_eq!(
+            azblob_builder.account_key.unwrap(),
+            "account-key".to_string()
+        );
+
+        assert_eq!(azblob_builder.sas_token.unwrap(), "sas".to_string());
+    }
+
     #[test]
     fn test_builder_from_connection_string() {
         let builder = AzblobBuilder::from_connection_string(

src/services/azdfs/backend.rs

@@ -16,17 +16,18 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 use std::fmt::Formatter;
 use std::fmt::Write;
-use std::mem;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use ctor::ctor;
 use http::header::CONTENT_DISPOSITION;
 use http::header::CONTENT_LENGTH;
 use http::header::CONTENT_TYPE;
 use http::Request;
 use http::Response;
 use http::StatusCode;
 use log::debug;
+use regex::Regex;
 use reqsign::AzureStorageSigner;
 
 use super::error::parse_error;
@@ -36,6 +37,15 @@ use crate::ops::*;
 use crate::raw::*;
 use crate::*;
 
+/// Known regex pattern Azure Data Lake Storage Gen2 URI syntax.
+/// Azure public cloud: https://accountname.dfs.core.windows.net
+/// Azure US Government: https://accountname.dfs.core.usgovcloudapi.net
+/// Azure China: https://accountname.dfs.core.chinacloudapi.cn
+#[ctor]
+static KNOWN_AZDFS_RESOURCE_URI_SYNTAX_REGEX: Regex =
+    Regex::new(r"(?i)^https?://([a-z0-9]{3,24})\.dfs\.core\.(?:usgovcloudapi\.net|chinacloudapi\.cn|windows\.net)/?$")
+    .unwrap();
+
 /// Azure Data Lake Storage Gen2 Support.
 ///
 /// As known as `abfs`, `azdfs` or `azdls`.
@@ -241,8 +251,15 @@ impl Builder for AzdfsBuilder {
         };
 
         let mut signer_builder = AzureStorageSigner::builder();
+        let mut account_name = None;
         if let (Some(name), Some(key)) = (&self.account_name, &self.account_key) {
+            account_name = Some(name.clone());
             signer_builder.account_name(name).account_key(key);
+        } else if let Some(key) = &self.account_key {
+            account_name = infer_storage_name_from_endpoint(endpoint.as_str());
+            signer_builder
+                .account_name(account_name.as_ref().unwrap_or(&String::new()))
+                .account_key(key);
         }
 
         let signer = signer_builder.build().map_err(|e| {
@@ -261,7 +278,7 @@ impl Builder for AzdfsBuilder {
             signer: Arc::new(signer),
             filesystem: self.filesystem.clone(),
             client,
-            _account_name: mem::take(&mut self.account_name).unwrap_or_default(),
+            _account_name: account_name.unwrap_or_else(String::new),
         })
     }
 
@@ -597,3 +614,83 @@ impl AzdfsBackend {
         self.client.send_async(req).await
     }
 }
+
+fn infer_storage_name_from_endpoint(endpoint: &str) -> Option<String> {
+    let candidate_names = KNOWN_AZDFS_RESOURCE_URI_SYNTAX_REGEX.captures(endpoint);
+    if candidate_names.is_none() || candidate_names.as_ref().unwrap().len() != 2 {
+        return None;
+    }
+    Some(
+        candidate_names
+            .unwrap()
+            .get(1)
+            .unwrap()
+            .clone()
+            .as_str()
+            .to_string(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{services::azdfs::backend::infer_storage_name_from_endpoint, Builder};
+
+    use super::AzdfsBuilder;
+
+    #[test]
+    fn test_infer_storage_name_from_endpoint() {
+        let endpoint = "https://account.dfs.core.windows.net";
+        let storage_name = infer_storage_name_from_endpoint(endpoint);
+        assert_eq!(storage_name, Some("account".to_string()));
+    }
+
+    #[test]
+    fn test_infer_storage_name_from_endpoint_with_trailing_slash() {
+        let endpoint = "https://account.dfs.core.windows.net/";
+        let storage_name = infer_storage_name_from_endpoint(endpoint);
+        assert_eq!(storage_name, Some("account".to_string()));
+    }
+
+    #[test]
+    fn test_builder_from_endpoint_and_key_infer_account_name() {
+        let mut azdfs_builder = AzdfsBuilder::default();
+        azdfs_builder.endpoint("https://storagesample.dfs.core.chinacloudapi.cn");
+        azdfs_builder.account_key("account-key");
+        azdfs_builder.filesystem("filesystem");
+        let azdfs = azdfs_builder
+            .build()
+            .expect("build azdfs should be succeeded.");
+
+        assert_eq!(
+            azdfs.endpoint,
+            "https://storagesample.dfs.core.chinacloudapi.cn"
+        );
+
+        assert_eq!(azdfs._account_name, "storagesample".to_string());
+
+        assert_eq!(azdfs.filesystem, "filesystem".to_string());
+
+        assert_eq!(
+            azdfs_builder.account_key.unwrap(),
+            "account-key".to_string()
+        );
+    }
+
+    #[test]
+    fn test_no_key_wont_infer_account_name() {
+        let mut azdfs_builder = AzdfsBuilder::default();
+        azdfs_builder.endpoint("https://storagesample.dfs.core.windows.net");
+        azdfs_builder.filesystem("filesystem");
+        let azdfs = azdfs_builder
+            .build()
+            .expect("build azdfs should be succeeded.");
+
+        assert_eq!(azdfs.endpoint, "https://storagesample.dfs.core.windows.net");
+
+        assert_eq!(azdfs._account_name, "".to_string());
+
+        assert_eq!(azdfs.filesystem, "filesystem".to_string());
+
+        assert_eq!(azdfs_builder.account_key, None);
+    }
+}