From 05b64c6be9c22dbf608c9fd90bea85799b47d478 Mon Sep 17 00:00:00 2001 From: Tomas Sedlak Date: Fri, 7 Apr 2023 22:45:04 +0200 Subject: [PATCH] Simplify the Store Backend Configuration code --- rust/src/storage/config.rs | 78 +++++++++++++------------------------- rust/src/storage/mod.rs | 12 ++---- 2 files changed, 29 insertions(+), 61 deletions(-) diff --git a/rust/src/storage/config.rs b/rust/src/storage/config.rs index e447bc687a..682b104105 100644 --- a/rust/src/storage/config.rs +++ b/rust/src/storage/config.rs @@ -5,7 +5,7 @@ use crate::{DeltaResult, DeltaTableError}; use object_store::memory::InMemory; use object_store::path::Path; use object_store::prefix::PrefixStore; -use object_store::DynObjectStore; +use object_store::{DynObjectStore, ObjectStore}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; @@ -16,9 +16,9 @@ use super::s3::{S3StorageBackend, S3StorageOptions}; #[cfg(any(feature = "s3", feature = "s3-native-tls"))] use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; #[cfg(feature = "azure")] -use object_store::azure::{AzureConfigKey, MicrosoftAzure, MicrosoftAzureBuilder}; +use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder}; #[cfg(feature = "gcs")] -use object_store::gcp::{GoogleCloudStorage, GoogleCloudStorageBuilder, GoogleConfigKey}; +use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey}; #[cfg(any( feature = "s3", feature = "s3-native-tls", @@ -97,45 +97,6 @@ impl From> for StorageOptions { } } -pub(crate) enum ObjectStoreImpl { - Local(FileStorageBackend), - InMemory(InMemory), - #[cfg(any(feature = "s3", feature = "s3-native-tls"))] - S3(S3StorageBackend), - #[cfg(feature = "gcs")] - Google(GoogleCloudStorage), - #[cfg(feature = "azure")] - Azure(MicrosoftAzure), -} - -impl ObjectStoreImpl { - pub(crate) fn into_prefix(self, prefix: Path) -> Arc { - match self { - ObjectStoreImpl::Local(store) => Arc::new(PrefixStore::new(store, prefix)), - ObjectStoreImpl::InMemory(store) => Arc::new(PrefixStore::new(store, prefix)), - #[cfg(feature = "azure")] - ObjectStoreImpl::Azure(store) => Arc::new(PrefixStore::new(store, prefix)), - #[cfg(any(feature = "s3", feature = "s3-native-tls"))] - ObjectStoreImpl::S3(store) => Arc::new(PrefixStore::new(store, prefix)), - #[cfg(feature = "gcs")] - ObjectStoreImpl::Google(store) => Arc::new(PrefixStore::new(store, prefix)), - } - } - - pub(crate) fn into_store(self) -> Arc { - match self { - ObjectStoreImpl::Local(store) => Arc::new(store), - ObjectStoreImpl::InMemory(store) => Arc::new(store), - #[cfg(feature = "azure")] - ObjectStoreImpl::Azure(store) => Arc::new(store), - #[cfg(any(feature = "s3", feature = "s3-native-tls"))] - ObjectStoreImpl::S3(store) => Arc::new(store), - #[cfg(feature = "gcs")] - ObjectStoreImpl::Google(store) => Arc::new(store), - } - } -} - pub(crate) enum ObjectStoreKind { Local, InMemory, @@ -176,24 +137,28 @@ impl ObjectStoreKind { pub fn into_impl( self, - storage_url: impl AsRef, + storage_url: &Url, options: impl Into, - ) -> DeltaResult { + ) -> DeltaResult> { let _options = options.into(); match self { - ObjectStoreKind::Local => Ok(ObjectStoreImpl::Local(FileStorageBackend::new())), - ObjectStoreKind::InMemory => Ok(ObjectStoreImpl::InMemory(InMemory::new())), + ObjectStoreKind::Local => Ok(Self::url_prefix_handler( + FileStorageBackend::new(), + storage_url, + )), + ObjectStoreKind::InMemory => Ok(Self::url_prefix_handler(InMemory::new(), storage_url)), #[cfg(any(feature = "s3", feature = "s3-native-tls"))] ObjectStoreKind::S3 => { - let store = AmazonS3Builder::from_env() + let amazon_s3 = AmazonS3Builder::from_env() .with_url(storage_url.as_ref()) .try_with_options(&_options.as_s3_options())? .with_allow_http(_options.allow_http()) .build()?; - Ok(ObjectStoreImpl::S3(S3StorageBackend::try_new( - Arc::new(store), + let store = S3StorageBackend::try_new( + Arc::new(amazon_s3), S3StorageOptions::from_map(&_options.0), - )?)) + )?; + Ok(Self::url_prefix_handler(store, storage_url)) } #[cfg(not(any(feature = "s3", feature = "s3-native-tls")))] ObjectStoreKind::S3 => Err(DeltaTableError::MissingFeature { @@ -207,7 +172,7 @@ impl ObjectStoreKind { .try_with_options(&_options.as_azure_options())? .with_allow_http(_options.allow_http()) .build()?; - Ok(ObjectStoreImpl::Azure(store)) + Ok(Self::url_prefix_handler(store, storage_url)) } #[cfg(not(feature = "azure"))] ObjectStoreKind::Azure => Err(DeltaTableError::MissingFeature { @@ -220,7 +185,7 @@ impl ObjectStoreKind { .with_url(storage_url.as_ref()) .try_with_options(&_options.as_gcs_options())? .build()?; - Ok(ObjectStoreImpl::Google(store)) + Ok(Self::url_prefix_handler(store, storage_url)) } #[cfg(not(feature = "gcs"))] ObjectStoreKind::Google => Err(DeltaTableError::MissingFeature { @@ -229,4 +194,13 @@ impl ObjectStoreKind { }), } } + + fn url_prefix_handler(store: T, storage_url: &Url) -> Arc { + let prefix = Path::from(storage_url.path()); + if prefix != Path::from("/") { + Arc::new(PrefixStore::new(store, prefix)) + } else { + Arc::new(store) + } + } } diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index 19e955a958..c0768bc9d8 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -89,14 +89,8 @@ impl DeltaObjectStore { /// * `location` - A url pointing to the root of the delta table. /// * `options` - Options passed to underlying builders. See [`with_storage_options`](crate::builder::DeltaTableBuilder::with_storage_options) pub fn try_new(location: Url, options: impl Into + Clone) -> DeltaResult { - let prefix = Path::from(location.path()); - let root_store = - ObjectStoreKind::parse_url(&location)?.into_impl(location.as_ref(), options.clone())?; - let storage = if prefix != Path::from("/") { - root_store.into_prefix(prefix) - } else { - root_store.into_store() - }; + let storage = + ObjectStoreKind::parse_url(&location)?.into_impl(&location, options.clone())?; Ok(Self { storage, location, @@ -109,7 +103,7 @@ impl DeltaObjectStore { self.storage.clone() } - /// Storage options used to intialize storage backend + /// Storage options used to initialize storage backend pub fn storage_options(&self) -> &StorageOptions { &self.options }