feat: centrally apply object store layers

Signed-off-by: Robert Pack <[email protected]>

Author: roeap

crates/aws/src/lib.rs

@@ -27,7 +27,7 @@ use aws_sdk_dynamodb::{
 };
 use deltalake_core::logstore::object_store::aws::AmazonS3ConfigKey;
 use deltalake_core::logstore::{default_logstore, logstores, LogStore, LogStoreFactory};
-use deltalake_core::logstore::{factories, url_prefix_handler, ObjectStoreRef, StorageConfig};
+use deltalake_core::logstore::{factories, ObjectStoreRef, StorageConfig};
 use deltalake_core::{DeltaResult, Path};
 use errors::{DynamoDbConfigError, LockClientError};
 use regex::Regex;
@@ -55,7 +55,6 @@ impl LogStoreFactory for S3LogStoreFactory {
         location: &Url,
         options: &StorageConfig,
     ) -> DeltaResult<Arc<dyn LogStore>> {
-        let store = url_prefix_handler(store, Path::parse(location.path())?);
         let s3_options = self.with_env_s3(&options.raw.clone().into());
         if s3_options.keys().any(|key| {
             let key = key.to_ascii_lowercase();

crates/aws/src/storage.rs

@@ -1,4 +1,10 @@
 //! AWS S3 storage backend.
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::ops::Range;
+use std::str::FromStr;
+use std::sync::Arc;
+use std::time::Duration;
 
 use aws_config::{Region, SdkConfig};
 use bytes::Bytes;
@@ -7,18 +13,12 @@ use deltalake_core::logstore::object_store::{
     GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreScheme,
     PutMultipartOpts, PutOptions, PutPayload, PutResult, Result as ObjectStoreResult,
 };
-use deltalake_core::logstore::{
-    limit_store_handler, str_is_truthy, ObjectStoreFactory, ObjectStoreRef, StorageConfig,
-};
+use deltalake_core::logstore::{str_is_truthy, ObjectStoreFactory, ObjectStoreRef};
 use deltalake_core::{DeltaResult, DeltaTableError, ObjectStoreError, Path};
 use futures::stream::BoxStream;
 use futures::Future;
-use std::collections::HashMap;
-use std::fmt::Debug;
-use std::ops::Range;
-use std::str::FromStr;
-use std::sync::Arc;
-use std::time::Duration;
+use object_store::aws::AmazonS3;
+use object_store::RetryConfig;
 use tracing::log::*;
 use url::Url;
 
@@ -36,9 +36,10 @@ impl ObjectStoreFactory for S3ObjectStoreFactory {
     fn parse_url_opts(
         &self,
         url: &Url,
-        storage_options: &StorageConfig,
+        storage_options: &HashMap<String, String>,
+        retry: &RetryConfig,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
-        let options = self.with_env_s3(&storage_options.raw);
+        let options = self.with_env_s3(storage_options);
 
         // All S3-likes should start their builder the same way
         let mut builder = AmazonS3Builder::new().with_url(url.to_string());
@@ -63,31 +64,30 @@ impl ObjectStoreFactory for S3ObjectStoreFactory {
             ));
         }
 
-        let inner = builder.with_retry(storage_options.retry.clone()).build()?;
-        let limit = storage_options.limit.clone().unwrap_or_default();
-        let store = aws_storage_handler(limit_store_handler(inner, &limit), &s3_options)?;
+        let inner = builder.with_retry(retry.clone()).build()?;
+        let store = aws_storage_handler(inner, &s3_options)?;
         debug!("Initialized the object store: {store:?}");
 
         Ok((store, prefix))
     }
 }
 
 fn aws_storage_handler(
-    store: ObjectStoreRef,
+    store: AmazonS3,
     s3_options: &S3StorageOptions,
 ) -> DeltaResult<ObjectStoreRef> {
     // Nearly all S3 Object stores support conditional put, so we change the default to always returning an S3 Object store
     // unless explicitly passing a locking provider key or allow_unsafe_rename. Then we will pass it to the old S3StorageBackend.
     if s3_options.locking_provider.as_deref() == Some("dynamodb") || s3_options.allow_unsafe_rename
     {
         let store = S3StorageBackend::try_new(
-            store,
+            Arc::new(store),
             Some("dynamodb") == s3_options.locking_provider.as_deref()
                 || s3_options.allow_unsafe_rename,
         )?;
         Ok(Arc::new(store))
     } else {
-        Ok(store)
+        Ok(Arc::new(store))
     }
 }
 

crates/aws/tests/integration_s3_dynamodb.rs

@@ -165,7 +165,7 @@ async fn test_repair_commit_entry() -> TestResult<()> {
     let options: StorageConfig = OPTIONS.clone().into_iter().collect();
     let log_store: S3DynamoDbLogStore = S3DynamoDbLogStore::try_new(
         ensure_table_uri(table.table_uri())?,
-        options.clone(),
+        &options,
         &S3_OPTIONS,
         std::sync::Arc::new(table.object_store()),
     )?;
@@ -240,7 +240,7 @@ async fn test_abort_commit_entry() -> TestResult<()> {
     let options: StorageConfig = OPTIONS.clone().into_iter().collect();
     let log_store: S3DynamoDbLogStore = S3DynamoDbLogStore::try_new(
         ensure_table_uri(table.table_uri())?,
-        options.clone(),
+        &options,
         &S3_OPTIONS,
         std::sync::Arc::new(table.object_store()),
     )?;
@@ -287,7 +287,7 @@ async fn test_abort_commit_entry_fail_to_delete_entry() -> TestResult<()> {
     let options: StorageConfig = OPTIONS.clone().into_iter().collect();
     let log_store: S3DynamoDbLogStore = S3DynamoDbLogStore::try_new(
         ensure_table_uri(table.table_uri())?,
-        options.clone(),
+        &options,
         &S3_OPTIONS,
         std::sync::Arc::new(table.object_store()),
     )?;

crates/azure/src/lib.rs

@@ -3,13 +3,10 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use deltalake_core::logstore::{default_logstore, logstores, LogStore, LogStoreFactory};
-use deltalake_core::logstore::{
-    factories, limit_store_handler, url_prefix_handler, ObjectStoreFactory, ObjectStoreRef,
-    StorageConfig,
-};
+use deltalake_core::logstore::{factories, ObjectStoreFactory, ObjectStoreRef, StorageConfig};
 use deltalake_core::{DeltaResult, DeltaTableError, Path};
 use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder};
-use object_store::ObjectStoreScheme;
+use object_store::{ObjectStoreScheme, RetryConfig};
 use url::Url;
 
 mod config;
@@ -19,10 +16,9 @@ trait AzureOptions {
     fn as_azure_options(&self) -> HashMap<AzureConfigKey, String>;
 }
 
-impl AzureOptions for StorageConfig {
+impl AzureOptions for HashMap<String, String> {
     fn as_azure_options(&self) -> HashMap<AzureConfigKey, String> {
-        self.raw
-            .iter()
+        self.iter()
             .filter_map(|(key, value)| {
                 Some((
                     AzureConfigKey::from_str(&key.to_ascii_lowercase()).ok()?,
@@ -40,7 +36,8 @@ impl ObjectStoreFactory for AzureFactory {
     fn parse_url_opts(
         &self,
         url: &Url,
-        options: &StorageConfig,
+        options: &HashMap<String, String>,
+        retry: &RetryConfig,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
         let config = config::AzureConfigHelper::try_new(options.as_azure_options())?.build()?;
 
@@ -51,15 +48,12 @@ impl ObjectStoreFactory for AzureFactory {
         let prefix = Path::parse(path)?;
 
         let mut builder = MicrosoftAzureBuilder::new().with_url(url.to_string());
-
         for (key, value) in config.iter() {
             builder = builder.with_config(*key, value.clone());
         }
+        let store = builder.with_retry(retry.clone()).build()?;
 
-        let inner = builder.with_retry(options.retry.clone()).build()?;
-        let limit = options.limit.clone().unwrap_or_default();
-        let store = limit_store_handler(url_prefix_handler(inner, prefix.clone()), &limit);
-        Ok((store, prefix))
+        Ok((Arc::new(store), prefix))
     }
 }
 

crates/catalog-unity/src/lib.rs

@@ -9,7 +9,8 @@ compile_error!(
 
 use datafusion_common::DataFusionError;
 use deltalake_core::logstore::{
-    default_logstore, logstores, LogStore, LogStoreFactory, StorageConfig,
+    default_logstore, logstores, object_store::RetryConfig, LogStore, LogStoreFactory,
+    StorageConfig,
 };
 use reqwest::header::{HeaderValue, InvalidHeaderValue, AUTHORIZATION};
 use reqwest::Url;
@@ -840,23 +841,25 @@ impl ObjectStoreFactory for UnityCatalogFactory {
     fn parse_url_opts(
         &self,
         table_uri: &Url,
-        options: &StorageConfig,
+        options: &HashMap<String, String>,
+        _retry: &RetryConfig,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
         let (table_path, temp_creds) = UnityCatalogBuilder::execute_uc_future(
             UnityCatalogBuilder::get_uc_location_and_token(table_uri.as_str()),
         )??;
 
-        let mut storage_options = options.raw.clone();
+        let mut storage_options = options.clone();
         storage_options.extend(temp_creds);
 
+        // TODO(roeap): we should not have to go through the table here.
+        // ideally we just create the right storage ...
         let mut builder =
             DeltaTableBuilder::from_uri(&table_path).with_io_runtime(IORuntime::default());
         if !storage_options.is_empty() {
             builder = builder.with_storage_options(storage_options.clone());
         }
-
         let prefix = Path::parse(table_uri.path())?;
-        let store = builder.build()?.object_store();
+        let store = builder.build_storage()?.object_store(None);
 
         Ok((store, prefix))
     }

crates/core/Cargo.toml

@@ -30,6 +30,7 @@ arrow-row = { workspace = true }
 arrow-schema = { workspace = true, features = ["serde"] }
 arrow-select = { workspace = true }
 parquet = { workspace = true, features = ["async", "object_store"] }
+object_store = { workspace = true, features = ["cloud"] }
 pin-project-lite = "^0.2.7"
 
 # datafusion
@@ -87,7 +88,6 @@ itertools = "0.14"
 libc = ">=0.2.90, <1"
 num-bigint = "0.4"
 num-traits = "0.2.15"
-object_store = { workspace = true }
 parking_lot = "0.12"
 percent-encoding = "2"
 roaring = "0.10.1"

crates/core/src/data_catalog/storage/mod.rs

@@ -49,7 +49,6 @@ impl ListingSchemaProvider {
     ) -> DeltaResult<Self> {
         let uri = ensure_table_uri(root_uri)?;
         let options = options.unwrap_or_default();
-        // We already parsed the url, so unwrapping is safe.
         let store = store_for(&uri, &options)?;
         Ok(Self {
             authority: uri.to_string(),

crates/core/src/logstore/config.rs

@@ -1,7 +1,8 @@
 use std::collections::HashMap;
 
-#[cfg(feature = "cloud")]
 use ::object_store::RetryConfig;
+use object_store::{path::Path, prefix::PrefixStore, ObjectStore};
+use tokio::runtime::Handle;
 
 #[cfg(feature = "delta-cache")]
 use super::storage::cache::LogCacheConfig;
@@ -104,16 +105,56 @@ pub struct StorageConfig {
     /// Configuration to set up a dedicated IO runtime to execute IO related operations.
     pub runtime: Option<RuntimeConfig>,
 
-    #[cfg(feature = "cloud")]
     pub retry: ::object_store::RetryConfig,
 
+    /// Limit configuration.
+    ///
+    /// Configuration to limit the number of concurrent requests to the object store.
     pub limit: Option<LimitConfig>,
 
     pub unknown_properties: HashMap<String, String>,
 
     pub raw: HashMap<String, String>,
 }
 
+impl StorageConfig {
+    /// Wrap an object store with additional layers of functionality.
+    ///
+    /// Depending on the configuration, the following layers may be added:
+    /// - Retry layer: Adds retry logic to the object store.
+    /// - Limit layer: Limits the number of concurrent requests to the object store.
+    /// - Runtime layer: Executes IO related operations on a dedicated runtime.
+    pub fn decorate_store<T: ObjectStore + Clone>(
+        &self,
+        store: T,
+        table_root: &url::Url,
+        handle: Option<Handle>,
+    ) -> DeltaResult<Box<dyn ObjectStore>> {
+        let inner = if let Some(runtime) = &self.runtime {
+            Box::new(runtime.decorate(store, handle)) as Box<dyn ObjectStore>
+        } else {
+            Box::new(store) as Box<dyn ObjectStore>
+        };
+
+        let prefix = if table_root.scheme() == "file" {
+            Path::from_filesystem_path(
+                table_root
+                    .to_file_path()
+                    .map_err(|_| DeltaTableError::generic("failed to convert fs"))?,
+            )?
+        } else {
+            Path::from_url_path(table_root.path())?
+        };
+        let inner = if prefix != Path::from("/") {
+            Box::new(PrefixStore::new(inner, prefix)) as Box<dyn ObjectStore>
+        } else {
+            Box::new(inner) as Box<dyn ObjectStore>
+        };
+
+        Ok(inner)
+    }
+}
+
 impl<K, V> FromIterator<(K, V)> for StorageConfig
 where
     K: AsRef<str> + Into<String>,