feat: extend configuration system

Author: roeap

Cargo.lock

@@ -2,4 +2,4 @@
 members = ["object-store"]
 
 [patch.crates-io]
-object_store = { git = "https://github.com/apache/arrow-rs", rev = "2408bb274e82c785cca9b4596cd8f201ccc5d7c6" }
+object_store = { git = "https://github.com/apache/arrow-rs", rev = "1889e33da31218ee2c58ad874036b17b699538b9" }

Cargo.toml

@@ -10,6 +10,7 @@ bytes = "1.2.1"
 futures = "0.3"
 once_cell = "1.12.0"
 object_store = { version = "0.5.2", features = ["azure", "aws", "gcp"] }
+percent-encoding = "2"
 pyo3 = { version = "0.17.1", features = [
   "extension-module",
   "abi3",

object-store/Cargo.toml

@@ -8,35 +8,96 @@ use object_store::local::LocalFileSystem;
 use object_store::memory::InMemory;
 use object_store::path::Path;
 use object_store::prefix::PrefixObjectStore;
-use object_store::{ClientOptions, DynObjectStore, Result as ObjectStoreResult};
+use object_store::{
+    ClientOptions, DynObjectStore, Error as ObjectStoreError, Result as ObjectStoreResult,
+    RetryConfig,
+};
 use url::Url;
 
-use crate::settings::AzureConfig;
+use crate::config::{aws::S3Config, azure::AzureConfig, google::GoogleConfig};
 
-pub enum ObjectStoreImpl {
+enum ObjectStoreKind {
+    Local,
+    InMemory,
+    S3,
+    Google,
+    Azure,
+}
+
+impl ObjectStoreKind {
+    pub fn parse_url(url: &Url) -> ObjectStoreResult<Self> {
+        match url.scheme() {
+            "file" => Ok(ObjectStoreKind::Local),
+            "memory" => Ok(ObjectStoreKind::InMemory),
+            "az" | "abfs" | "abfss" | "azure" | "wasb" | "adl" => Ok(ObjectStoreKind::Azure),
+            "s3" | "s3a" => Ok(ObjectStoreKind::S3),
+            "gs" => Ok(ObjectStoreKind::Google),
+            "https" => {
+                let host = url.host_str().unwrap_or_default();
+                if host.contains("amazonaws.com") {
+                    Ok(ObjectStoreKind::S3)
+                } else if host.contains("dfs.core.windows.net")
+                    || host.contains("blob.core.windows.net")
+                {
+                    Ok(ObjectStoreKind::Azure)
+                } else {
+                    Err(ObjectStoreError::NotImplemented)
+                }
+            }
+            _ => Err(ObjectStoreError::NotImplemented),
+        }
+    }
+}
+
+enum ObjectStoreImpl {
     Local(LocalFileSystem),
+    InMemory(InMemory),
     Azrue(MicrosoftAzure),
     S3(AmazonS3),
     Gcp(GoogleCloudStorage),
-    InMemory(InMemory),
 }
 
-pub struct StorageBuilder {
+impl ObjectStoreImpl {
+    pub fn into_prefix(self, prefix: Path) -> Arc<DynObjectStore> {
+        match self {
+            ObjectStoreImpl::Local(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
+            ObjectStoreImpl::InMemory(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
+            ObjectStoreImpl::Azrue(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
+            ObjectStoreImpl::S3(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
+            ObjectStoreImpl::Gcp(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
+        }
+    }
+
+    pub fn into_store(self) -> Arc<DynObjectStore> {
+        match self {
+            ObjectStoreImpl::Local(store) => Arc::new(store),
+            ObjectStoreImpl::InMemory(store) => Arc::new(store),
+            ObjectStoreImpl::Azrue(store) => Arc::new(store),
+            ObjectStoreImpl::S3(store) => Arc::new(store),
+            ObjectStoreImpl::Gcp(store) => Arc::new(store),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ObjectStoreBuilder {
     url: String,
     prefix: Option<Path>,
     path_as_prefix: bool,
     options: HashMap<String, String>,
-    client_options: ClientOptions,
+    client_options: Option<ClientOptions>,
+    retry_config: Option<RetryConfig>,
 }
 
-impl StorageBuilder {
+impl ObjectStoreBuilder {
     pub fn new(url: impl Into<String>) -> Self {
         Self {
             url: url.into(),
             prefix: None,
             path_as_prefix: false,
             options: Default::default(),
-            client_options: Default::default(),
+            client_options: None,
+            retry_config: None,
         }
     }
 
@@ -55,39 +116,60 @@ impl StorageBuilder {
         self
     }
 
+    pub fn with_path_as_prefix(mut self, path_as_prefix: bool) -> Self {
+        self.path_as_prefix = path_as_prefix;
+        self
+    }
+
+    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
+        self.client_options = Some(options);
+        self
+    }
+
+    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
+        self.retry_config = Some(retry_config);
+        self
+    }
+
     pub fn build(mut self) -> ObjectStoreResult<Arc<DynObjectStore>> {
-        let url = Url::parse(&self.url).unwrap();
-        let root_store = match url.scheme() {
-            "file" => ObjectStoreImpl::Local(LocalFileSystem::new()),
-            "memory" => ObjectStoreImpl::InMemory(InMemory::new()),
-            "az" | "abfs" | "abfss" | "azure" | "wasb" | "adl" => {
-                let mut builder = AzureConfig::get_builder(&url, &self.options)?;
-                builder = builder.with_client_options(self.client_options);
+        let url = Url::parse(&self.url).map_err(|err| ObjectStoreError::Generic {
+            store: "Generic",
+            source: Box::new(err),
+        })?;
+        let root_store = match ObjectStoreKind::parse_url(&url).unwrap() {
+            ObjectStoreKind::Local => ObjectStoreImpl::Local(LocalFileSystem::new()),
+            ObjectStoreKind::InMemory => ObjectStoreImpl::InMemory(InMemory::new()),
+            ObjectStoreKind::Azure => {
+                let mut builder = AzureConfig::get_builder(&self.url, &self.options)?;
+                builder = builder
+                    .with_client_options(self.client_options.unwrap_or_default())
+                    .with_retry(self.retry_config.unwrap_or_default());
                 ObjectStoreImpl::Azrue(builder.build()?)
             }
-            "s3" | "s3a" => todo!(),
-            "gs" => todo!(),
-            "https" => todo!(),
-            _ => todo!(),
+            ObjectStoreKind::S3 => {
+                let mut builder = S3Config::get_builder(&self.url, &self.options)?;
+                builder = builder
+                    .with_client_options(self.client_options.unwrap_or_default())
+                    .with_retry(self.retry_config.unwrap_or_default());
+                ObjectStoreImpl::S3(builder.build()?)
+            }
+            ObjectStoreKind::Google => {
+                let mut builder = GoogleConfig::get_builder(&self.url, &self.options)?;
+                builder = builder
+                    .with_client_options(self.client_options.unwrap_or_default())
+                    .with_retry(self.retry_config.unwrap_or_default());
+                ObjectStoreImpl::Gcp(builder.build()?)
+            }
         };
 
         if self.path_as_prefix && !url.path().is_empty() && self.prefix.is_none() {
             self.prefix = Some(Path::from(url.path()))
         }
 
-        let store: Arc<DynObjectStore> = if let Some(prefix) = self.prefix {
-            match root_store {
-                ObjectStoreImpl::Local(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
-                ObjectStoreImpl::InMemory(store) => Arc::new(PrefixObjectStore::new(store, prefix)),
-                _ => todo!(),
-            }
+        if let Some(prefix) = self.prefix {
+            Ok(root_store.into_prefix(prefix))
         } else {
-            match root_store {
-                ObjectStoreImpl::Local(store) => Arc::new(store),
-                ObjectStoreImpl::InMemory(store) => Arc::new(store),
-                _ => todo!(),
-            }
-        };
-        Ok(store)
+            Ok(root_store.into_store())
+        }
     }
 }